api.py

"""
FastAPI Backend for Knowledge Graph System
Provides REST API endpoints for graph creation and querying
"""

from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List, Dict
import os
import shutil
from datetime import datetime
import uuid
import json

from document_processor import DocumentProcessor
from text_preprocessor import TextPreprocessor
from entity_extractor import EntityExtractor
from relation_extractor import RelationExtractor
from graph_builder import KnowledgeGraphBuilder
from graph_visualizer import GraphVisualizer
from graph_querier import KnowledgeGraphQuerier
from semantic_retriever import SemanticRetriever
from config import OPENAI_API_KEY

# Initialize FastAPI app
app = FastAPI(
    title="Knowledge Graph API",
    description="API for creating and querying knowledge graphs from documents",
    version="1.0.0"
)

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

# Global storage for active knowledge graphs
active_graphs: Dict[str, Dict] = {}

# Global storage for chat history
chat_history: Dict[str, List[Dict]] = {}

# Pydantic models for request/response
class GraphCreationResponse(BaseModel):
    graph_id: str
    message: str
    entities_count: int
    relations_count: int
    statistics: Dict
    output_dir: str

class SemanticSearchRequest(BaseModel):
    graph_id: str
    query: str
    top_k: int = 5

class QuestionAnswerRequest(BaseModel):
    graph_id: str
    question: str

class EntityRelationsRequest(BaseModel):
    graph_id: str
    entity_name: str

class SPARQLQueryRequest(BaseModel):
    graph_id: str
    query: str

class ChatMessage(BaseModel):
    role: str
    content: str
    timestamp: str
    facts: Optional[List[str]] = None

class SaveChatRequest(BaseModel):
    graph_id: str
    messages: List[ChatMessage]


@app.get("/")
async def root():
    """Root endpoint"""
    return {
        "message": "Knowledge Graph API",
        "version": "1.0.0",
        "endpoints": {
            "health": "/health",
            "upload": "/upload",
            "graphs": "/graphs",
            "semantic_search": "/semantic_search",
            "question_answer": "/question_answer",
            "entity_relations": "/entity_relations",
            "sparql_query": "/sparql_query",
            "visualization": "/visualization/{graph_id}",
            "download_graph": "/download_graph/{graph_id}"
        }
    }


@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "active_graphs": len(active_graphs),
        "openai_configured": bool(OPENAI_API_KEY)
    }


@app.post("/upload", response_model=GraphCreationResponse)
async def upload_document(
    file: UploadFile = File(...),
    background_tasks: BackgroundTasks = None
):
    """
    Upload a document and create a knowledge graph
    Supports PDF, TXT, and DOCX files
    """
    
    # Validate file type
    allowed_extensions = {'.pdf', '.txt', '.docx', '.doc'}
    file_ext = os.path.splitext(file.filename)[1].lower()
    
    if file_ext not in allowed_extensions:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type: {file_ext}. Allowed: {allowed_extensions}"
        )
    
    # Generate unique ID
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    graph_id = f"{timestamp}_{uuid.uuid4().hex[:8]}"
    output_dir = f"output/kg_{graph_id}"
    os.makedirs(output_dir, exist_ok=True)
    
    # Save uploaded file temporarily
    temp_file_path = f"{output_dir}/uploaded_{file.filename}"
    
    try:
        with open(temp_file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # Process document
        doc_processor = DocumentProcessor()
        text = doc_processor.load_document(temp_file_path)
        
        upload_progress[graph_id].update({
            'progress': 20,
            'message': 'Document loaded, preprocessing text...',
            'logs': upload_progress[graph_id]['logs'] + [f"✓ Loaded document from {temp_file_path}"]
        })
        
        # Preprocess
        preprocessor = TextPreprocessor()
        preprocessed = preprocessor.preprocess(text)
        
        # Extract entities
        entity_extractor = EntityExtractor(openai_api_key=OPENAI_API_KEY)
        spacy_entities = entity_extractor.extract_entities_spacy(preprocessed['cleaned'])
        
        upload_progress[graph_id].update({
            'progress': 40,
            'message': f'Extracted {len(spacy_entities)} entities with spaCy, using LLM...',
            'logs': upload_progress[graph_id]['logs'] + [f"✓ Extracted {len(spacy_entities)} entities using spaCy"]
        })
        
        llm_entities = entity_extractor.extract_entities_llm(preprocessed['cleaned'])
        
        upload_progress[graph_id].update({
            'progress': 50,
            'message': 'Merging entities...',
            'logs': upload_progress[graph_id]['logs'] + [f"✓ Extracted {len(llm_entities)} entities using LLM"]
        })
        
        entities = entity_extractor.merge_entities(spacy_entities, llm_entities)
        
        upload_progress[graph_id].update({
            'progress': 60,
            'message': 'Extracting relations...',
            'logs': upload_progress[graph_id]['logs'] + [f"✓ Merged into {len(entities)} unique entities"]
        })
        
        # Extract relations
        relation_extractor = RelationExtractor(openai_api_key=OPENAI_API_KEY)
        pattern_relations = relation_extractor.extract_relations_pattern(
            preprocessed['cleaned'], entities
        )
        llm_relations = relation_extractor.extract_relations_llm(
            preprocessed['cleaned'], entities
        )
        relations = relation_extractor.merge_relations(pattern_relations, llm_relations)
        
        # Build knowledge graph
        kg_builder = KnowledgeGraphBuilder()
        kg_builder.build_from_extractions(entities, relations)
        
        upload_progress[graph_id].update({
            'progress': 85,
            'message': 'Saving graph...',
            'logs': upload_progress[graph_id]['logs'] + ["✓ Knowledge graph built"]
        })
        
        # Save graph
        kg_builder.save_rdf(f"{output_dir}/knowledge_graph.ttl")
        
        upload_progress[graph_id].update({
            'progress': 90,
            'message': 'Creating visualization...',
            'logs': upload_progress[graph_id]['logs'] + ["✓ Graph saved as RDF"]
        })
        
        # Create visualization
        visualizer = GraphVisualizer(kg_builder)
        visualizer.visualize(f"{output_dir}/knowledge_graph.png")
        
        upload_progress[graph_id].update({
            'progress': 95,
            'message': 'Indexing for search...',
            'logs': upload_progress[graph_id]['logs'] + ["✓ Visualization created"]
        })
        
        # Initialize semantic retriever
        retriever = SemanticRetriever(kg_builder)
        retriever.index_graph()
        
        # Initialize querier
        querier = KnowledgeGraphQuerier(kg_builder)
        
        # Get statistics
        stats = kg_builder.get_statistics()
        
        # Store in memory
        active_graphs[graph_id] = {
            'kg_builder': kg_builder,
            'retriever': retriever,
            'querier': querier,
            'entities': entities,
            'relations': relations,
            'output_dir': output_dir,
            'filename': file.filename,
            'created_at': datetime.now().isoformat(),
            'statistics': stats
        }
        
        # Save metadata
        metadata = {
            'graph_id': graph_id,
            'filename': file.filename,
            'created_at': datetime.now().isoformat(),
            'entities_count': len(entities),
            'relations_count': len(relations),
            'statistics': stats
        }
        
        with open(f"{output_dir}/metadata.json", 'w') as f:
            json.dump(metadata, f, indent=2)
        
        return GraphCreationResponse(
            graph_id=graph_id,
            message="Knowledge graph created successfully",
            entities_count=len(entities),
            relations_count=len(relations),
            statistics=stats,
            output_dir=output_dir
        )
        
    except Exception as e:
        # Cleanup on error
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        raise HTTPException(status_code=500, detail=str(e))
    
    finally:
        # Close uploaded file
        await file.close()


@app.get("/graphs")
async def list_graphs():
    """List all active knowledge graphs"""
    graphs_info = []
    
    for graph_id, data in active_graphs.items():
        graphs_info.append({
            'graph_id': graph_id,
            'filename': data.get('filename', 'Unknown'),
            'created_at': data['created_at'],
            'entities_count': len(data['entities']),
            'relations_count': len(data['relations']),
            'statistics': data['statistics']
        })
    
    return {"graphs": graphs_info, "total": len(graphs_info)}




@app.get("/chat_history/{graph_id}")
async def get_chat_history(graph_id: str):
    """Get chat history for a specific graph"""
    if graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    return {
        "graph_id": graph_id,
        "messages": chat_history.get(graph_id, []),
        "count": len(chat_history.get(graph_id, []))
    }


@app.post("/chat_history/{graph_id}")
async def save_chat_message(graph_id: str, message: ChatMessage):
    """Save a chat message to history"""
    if graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    if graph_id not in chat_history:
        chat_history[graph_id] = []
    
    chat_history[graph_id].append(message.dict())
    
    return {
        "message": "Chat message saved",
        "total_messages": len(chat_history[graph_id])
    }


@app.delete("/chat_history/{graph_id}")
async def clear_chat_history(graph_id: str):
    """Clear chat history for a specific graph"""
    if graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    if graph_id in chat_history:
        del chat_history[graph_id]
    
    return {"message": "Chat history cleared"}


@app.get("/graph/{graph_id}")
async def get_graph_info(graph_id: str):
    """Get information about a specific graph"""
    if graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    data = active_graphs[graph_id]
    
    return {
        'graph_id': graph_id,
        'filename': data['filename'],
        'created_at': data['created_at'],
        'entities_count': len(data['entities']),
        'relations_count': len(data['relations']),
        'statistics': data['statistics'],
        'sample_entities': [e['text'] for e in data['entities'][:10]]
    }


@app.post("/semantic_search")
async def semantic_search(request: SemanticSearchRequest):
    """Perform semantic search on knowledge graph"""
    if request.graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    retriever = active_graphs[request.graph_id]['retriever']
    
    try:
        results = retriever.search(request.query, top_k=request.top_k)
        return {"results": results, "query": request.query}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/question_answer")
async def question_answer(request: QuestionAnswerRequest):
    """Answer questions using LLM and knowledge graph"""
    if request.graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    if not OPENAI_API_KEY:
        raise HTTPException(
            status_code=400,
            detail="OpenAI API key not configured"
        )
    
    retriever = active_graphs[request.graph_id]['retriever']
    
    try:
        answer = retriever.answer_question_llm(request.question, OPENAI_API_KEY)
        
        # Also get relevant facts
        relevant_facts = retriever.search(request.question, top_k=5)
        
        return {
            "question": request.question,
            "answer": answer,
            "relevant_facts": relevant_facts
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/entity_relations")
async def entity_relations(request: EntityRelationsRequest):
    """Find all relations for a specific entity"""
    if request.graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    querier = active_graphs[request.graph_id]['querier']
    
    try:
        results = querier.find_entity_relations(request.entity_name)
        return {
            "entity": request.entity_name,
            "relations": results,
            "count": len(results)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/sparql_query")
async def sparql_query(request: SPARQLQueryRequest):
    """Execute custom SPARQL query"""
    if request.graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    querier = active_graphs[request.graph_id]['querier']
    
    try:
        results = querier.query(request.query)
        return {
            "results": results,
            "count": len(results)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Query execution failed: {str(e)}")


@app.get("/visualization/{graph_id}")
async def get_visualization(graph_id: str):
    """Get the knowledge graph visualization image"""
    if graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    output_dir = active_graphs[graph_id]['output_dir']
    image_path = f"{output_dir}/knowledge_graph.png"
    
    if not os.path.exists(image_path):
        raise HTTPException(status_code=404, detail="Visualization not found")
    
    return FileResponse(image_path, media_type="image/png")


@app.get("/download_graph/{graph_id}")
async def download_graph(graph_id: str):
    """Download the RDF knowledge graph file"""
    if graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    output_dir = active_graphs[graph_id]['output_dir']
    graph_path = f"{output_dir}/knowledge_graph.ttl"
    
    if not os.path.exists(graph_path):
        raise HTTPException(status_code=404, detail="Graph file not found")
    
    return FileResponse(
        graph_path,
        media_type="text/turtle",
        filename=f"knowledge_graph_{graph_id}.ttl"
    )


@app.get("/entities/{graph_id}")
async def get_entities(graph_id: str):
    """Get all entities from a graph"""
    if graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    entities = active_graphs[graph_id]['entities']
    
    return {
        "entities": entities,
        "count": len(entities)
    }


@app.delete("/graph/{graph_id}")
async def delete_graph(graph_id: str):
    """Delete a knowledge graph from memory"""
    if graph_id not in active_graphs:
        raise HTTPException(status_code=404, detail="Graph not found")
    
    # Remove from memory
    del active_graphs[graph_id]
    
    return {"message": f"Graph {graph_id} deleted successfully"}


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