PlantAI-rag/api/methodology_api.py

#!/usr/bin/env python3
"""
方法学提取 API - FastAPI 版本

提供接口：
- POST /extract - 从 PDF 提取方法学参数
- GET /result/{job_id} - 获取提取结果
"""

import os
import sys
import json
import uuid
from pathlib import Path
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, List, Optional

sys.path.insert(0, str(Path(__file__).parent.parent))

from src.rag.extractors.methodology.simple_extractor import SimpleExtractor

app = FastAPI(title="Methodology Extraction API")

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 存储提取结果
results = {}


class ExtractionResult(BaseModel):
    job_id: str
    status: str
    data: Optional[Dict] = None
    error: Optional[str] = None


@app.post("/extract")
async def extract_methodology(pdf: UploadFile = File(...)):
    """从 PDF 提取方法学参数"""
    
    # 生成 job_id
    job_id = str(uuid.uuid4())[:8]
    
    # 保存 PDF
    upload_dir = Path(__file__).parent.parent / 'uploads'
    upload_dir.mkdir(exist_ok=True)
    pdf_path = upload_dir / f"{job_id}_{pdf.filename}"
    
    with open(pdf_path, "wb") as f:
        content = await pdf.read()
        f.write(content)
    
    # 提取
    try:
        extractor = SimpleExtractor()
        result_text = extractor.extract(pdf_path)
        
        # 解析为结构化数据
        structured = parse_extraction_result(result_text)
        
        # 保存结果
        results[job_id] = {
            'status': 'completed',
            'pdf_path': str(pdf_path),
            'raw_text': result_text,
            'data': structured
        }
        
        return {
            'job_id': job_id,
            'status': 'completed',
            'data': structured
        }
        
    except Exception as e:
        results[job_id] = {
            'status': 'failed',
            'error': str(e)
        }
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/result/{job_id}")
async def get_result(job_id: str):
    """获取提取结果"""
    if job_id not in results:
        raise HTTPException(status_code=404, detail="Job not found")
    
    return results[job_id]


def parse_extraction_result(text: str) -> Dict:
    """解析提取结果为结构化数据"""
    
    data = {
        'species': {
            'target': '',
            'references': []
        },
        'gene_family': {
            'name': '',
            'pfam_id': ''
        },
        'stages': []
    }
    
    # 简单解析
    lines = text.split('\n')
    
    for line in lines:
        line = line.strip()
        
        # 物种
        if '目标物种' in line or 'target' in line.lower():
            if '：' in line:
                data['species']['target'] = line.split('：', 1)[-1].strip()
        elif '参考物种' in line or 'references' in line.lower():
            if '：' in line:
                refs = line.split('：', 1)[-1].strip()
                data['species']['references'] = [r.strip() for r in refs.split('、') if r.strip()]
        elif '基因家族' in line:
            if '：' in line:
                parts = line.split('：', 1)[-1].strip()
                if '(' in parts and 'Pfam:' in parts:
                    name = parts.split('(')[0].strip()
                    pfam = parts.split('Pfam:')[-1].split(')')[0].strip()
                    data['gene_family']['name'] = name
                    data['gene_family']['pfam_id'] = pfam
    
    # 默认阶段
    data['stages'] = [
        {
            'id': 'identify',
            'name': '1. 基因家族鉴定',
            'params': [
                {'name': 'tool', 'label': '工具', 'value': 'BLASTp + HMMER'},
                {'name': 'pfam_id', 'label': 'Pfam ID', 'value': 'PF03110'},
                {'name': 'e_value', 'label': 'E-value', 'value': ''},
                {'name': 'identity', 'label': 'Identity', 'value': ''}
            ]
        },
        {
            'id': 'phylogeny',
            'name': '2. 系统发育分析',
            'params': [
                {'name': 'software', 'label': '软件', 'value': 'MEGA 6.0'},
                {'name': 'method', 'label': '方法', 'value': 'Neighbor-Joining'},
                {'name': 'bootstrap', 'label': 'Bootstrap', 'value': '1000'}
            ]
        },
        {
            'id': 'motif',
            'name': '3. 保守基序分析',
            'params': [
                {'name': 'tool', 'label': '工具', 'value': 'MEME'},
                {'name': 'max_motifs', 'label': 'Max motifs', 'value': '10'}
            ]
        },
        {
            'id': 'structure',
            'name': '4. 基因结构分析',
            'params': [
                {'name': 'tool', 'label': '工具', 'value': 'GSDS'}
            ]
        },
        {
            'id': 'synteny',
            'name': '5. 共线性分析',
            'params': [
                {'name': 'tool', 'label': '工具', 'value': 'MCScanX'}
            ]
        },
        {
            'id': 'expression',
            'name': '6. 表达分析',
            'params': [
                {'name': 'method', 'label': '方法', 'value': 'qRT-PCR'}
            ]
        }
    ]
    
    return data


if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5001)