ComputerAssignments/CA10_knowledge_graphs/src/api/advanced_api.py

"""
Advanced API System for Knowledge Graph Operations
This module provides comprehensive API endpoints for Knowledge Graph operations
"""

import os
import json
import asyncio
from typing import Dict, List, Any, Optional, Union
from datetime import datetime
import uuid

# FastAPI and related
from fastapi import FastAPI, HTTPException, BackgroundTasks, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, FileResponse
from pydantic import BaseModel, Field
import uvicorn

# LangChain and LangGraph
from langchain_google_genai import ChatGoogleGenerativeAI

# Vector stores and embeddings
import chromadb
from sentence_transformers import SentenceTransformer

# Graph visualization
from pyvis.network import Network
import matplotlib.pyplot as plt
import seaborn as sns

import sys
import os

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from core.knowledge_graph import KnowledgeGraph, Triple, Entity, Relation
from core.scientific_kg import AdvancedScientificKG, ScientificEntity
from models.embeddings import TransE, DistMult, ComplEx, KGEmbeddingTrainer
from reasoning.langgraph_reasoning import (
    KnowledgeGraphReasoner,
    MultiAgentReasoningSystem,
)
from reasoning.rag_system import KnowledgeGraphRAG


# Pydantic models for API
class QueryRequest(BaseModel):
    query: str = Field(..., description="The question to ask")
    max_results: int = Field(
        default=5, ge=1, le=20, description="Maximum number of results"
    )
    include_reasoning: bool = Field(default=True, description="Include reasoning steps")
    use_multi_agent: bool = Field(
        default=False, description="Use multi-agent reasoning"
    )


class QueryResponse(BaseModel):
    query: str
    answer: str
    confidence: float
    reasoning_steps: List[str] = []
    entities_found: List[str] = []
    relations_found: List[str] = []
    execution_time: float
    timestamp: str


class TripleRequest(BaseModel):
    subject: str = Field(..., description="Subject entity")
    predicate: str = Field(..., description="Predicate/relation")
    object: str = Field(..., description="Object entity")


class EntityRequest(BaseModel):
    entity_id: str = Field(..., description="Entity identifier")
    name: str = Field(..., description="Entity name")
    entity_type: str = Field(default="unknown", description="Entity type")
    attributes: Dict[str, Any] = Field(
        default_factory=dict, description="Entity attributes"
    )


class RelationRequest(BaseModel):
    relation_id: str = Field(..., description="Relation identifier")
    name: str = Field(..., description="Relation name")
    domain: str = Field(default="entity", description="Domain type")
    range: str = Field(default="entity", description="Range type")


class KGStats(BaseModel):
    total_entities: int
    total_relations: int
    total_triples: int
    entity_types: Dict[str, int]
    relation_frequency: Dict[str, int]
    last_updated: str


class VisualizationRequest(BaseModel):
    max_entities: int = Field(default=20, ge=5, le=100)
    layout: str = Field(default="spring", description="Layout algorithm")
    save_path: Optional[str] = Field(
        default=None, description="Path to save visualization"
    )


class AdvancedKnowledgeGraphAPI:
    """Advanced API system for Knowledge Graph operations"""

    def __init__(self, kg: KnowledgeGraph):
        self.kg = kg
        self.app = FastAPI(
            title="Advanced Knowledge Graph API",
            description="Comprehensive API for Knowledge Graph operations with LangGraph and RAG",
            version="1.0.0",
        )

        # Initialize components
        self.reasoner = None
        self.multi_agent = None
        self.rag_system = None
        self.embedding_trainer = None

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

        # Setup routes
        self._setup_routes()

        # Initialize components
        self._initialize_components()

    def _initialize_components(self):
        """Initialize all components asynchronously"""
        try:
            # Initialize reasoning systems
            self.reasoner = KnowledgeGraphReasoner(self.kg)
            self.multi_agent = MultiAgentReasoningSystem(self.kg)

            # Initialize RAG system
            self.rag_system = KnowledgeGraphRAG(self.kg)

            # Initialize embedding trainer
            if len(self.kg.entities) > 0 and len(self.kg.relations) > 0:
                model = TransE(
                    len(self.kg.entities), len(self.kg.relations), embedding_dim=50
                )
                self.embedding_trainer = KGEmbeddingTrainer(model, self.kg)

            print("✅ All components initialized successfully")

        except Exception as e:
            print(f"⚠️ Some components failed to initialize: {e}")

    def _setup_routes(self):
        """Setup all API routes"""

        # Health check
        @self.app.get("/health")
        async def health_check():
            return {
                "status": "healthy",
                "timestamp": datetime.now().isoformat(),
                "kg_stats": {
                    "entities": len(self.kg.entities),
                    "relations": len(self.kg.relations),
                    "triples": len(self.kg.triples),
                },
            }

        # Knowledge Graph statistics
        @self.app.get("/stats", response_model=KGStats)
        async def get_stats():
            """Get comprehensive KG statistics"""

            # Entity type distribution
            entity_types = {}
            for entity in self.kg.entities.values():
                entity_type = entity.entity_type
                entity_types[entity_type] = entity_types.get(entity_type, 0) + 1

            # Relation frequency
            relation_freq = {}
            for triple in self.kg.triples:
                rel = triple.predicate
                relation_freq[rel] = relation_freq.get(rel, 0) + 1

            return KGStats(
                total_entities=len(self.kg.entities),
                total_relations=len(self.kg.relations),
                total_triples=len(self.kg.triples),
                entity_types=entity_types,
                relation_frequency=relation_freq,
                last_updated=datetime.now().isoformat(),
            )

        # Query endpoints
        @self.app.post("/query", response_model=QueryResponse)
        async def query_kg(request: QueryRequest):
            """Query the knowledge graph using advanced reasoning"""
            start_time = datetime.now()

            try:
                if request.use_multi_agent and self.multi_agent:
                    result = await self.multi_agent.reason(request.query)
                elif self.reasoner:
                    result = await self.reasoner.reason(request.query)
                else:
                    # Fallback to simple query
                    triples = self.kg.query_triples()
                    result = {
                        "query": request.query,
                        "answer": f"Found {len(triples)} triples in the knowledge graph.",
                        "confidence": 0.5,
                        "reasoning_steps": ["Simple query executed"],
                        "entities": [],
                        "evidence": [],
                    }

                execution_time = (datetime.now() - start_time).total_seconds()

                return QueryResponse(
                    query=request.query,
                    answer=result["answer"],
                    confidence=result.get("confidence", 0.5),
                    reasoning_steps=(
                        result.get("reasoning_steps", [])
                        if request.include_reasoning
                        else []
                    ),
                    entities_found=result.get("entities", []),
                    relations_found=result.get("evidence", []),
                    execution_time=execution_time,
                    timestamp=datetime.now().isoformat(),
                )

            except Exception as e:
                raise HTTPException(status_code=500, detail=f"Query failed: {str(e)}")

        # RAG query endpoint
        @self.app.post("/rag/query", response_model=QueryResponse)
        async def rag_query(request: QueryRequest):
            """Query using RAG system"""
            if not self.rag_system:
                raise HTTPException(status_code=503, detail="RAG system not available")

            start_time = datetime.now()

            try:
                result = await self.rag_system.query(
                    request.query, k=request.max_results
                )

                execution_time = (datetime.now() - start_time).total_seconds()

                return QueryResponse(
                    query=request.query,
                    answer=result.answer,
                    confidence=result.confidence,
                    reasoning_steps=(
                        result.reasoning_path if request.include_reasoning else []
                    ),
                    entities_found=result.entities_found,
                    relations_found=result.relations_found,
                    execution_time=execution_time,
                    timestamp=datetime.now().isoformat(),
                )

            except Exception as e:
                raise HTTPException(
                    status_code=500, detail=f"RAG query failed: {str(e)}"
                )

        # Entity operations
        @self.app.post("/entities")
        async def add_entity(entity: EntityRequest):
            """Add a new entity to the knowledge graph"""
            try:
                self.kg.add_entity(
                    entity.entity_id, entity.name, entity.entity_type, entity.attributes
                )
                return {"message": f"Entity {entity.entity_id} added successfully"}
            except Exception as e:
                raise HTTPException(status_code=400, detail=str(e))

        @self.app.get("/entities")
        async def get_entities():
            """Get all entities"""
            return {
                entity_id: {
                    "name": entity.name,
                    "type": entity.entity_type,
                    "attributes": entity.attributes,
                }
                for entity_id, entity in self.kg.entities.items()
            }

        @self.app.get("/entities/{entity_id}")
        async def get_entity(entity_id: str):
            """Get specific entity"""
            if entity_id not in self.kg.entities:
                raise HTTPException(status_code=404, detail="Entity not found")

            entity = self.kg.entities[entity_id]
            return {
                "entity_id": entity_id,
                "name": entity.name,
                "type": entity.entity_type,
                "attributes": entity.attributes,
            }

        # Relation operations
        @self.app.post("/relations")
        async def add_relation(relation: RelationRequest):
            """Add a new relation to the knowledge graph"""
            try:
                self.kg.add_relation(
                    relation.relation_id, relation.name, relation.domain, relation.range
                )
                return {
                    "message": f"Relation {relation.relation_id} added successfully"
                }
            except Exception as e:
                raise HTTPException(status_code=400, detail=str(e))

        @self.app.get("/relations")
        async def get_relations():
            """Get all relations"""
            return {
                rel_id: {
                    "name": relation.name,
                    "domain": relation.domain,
                    "range": relation.range,
                }
                for rel_id, relation in self.kg.relations.items()
            }

        # Triple operations
        @self.app.post("/triples")
        async def add_triple(triple: TripleRequest):
            """Add a new triple to the knowledge graph"""
            try:
                self.kg.add_triple(triple.subject, triple.predicate, triple.object)
                return {"message": "Triple added successfully"}
            except Exception as e:
                raise HTTPException(status_code=400, detail=str(e))

        @self.app.get("/triples")
        async def get_triples(
            subject: str = None, predicate: str = None, object: str = None
        ):
            """Query triples"""
            results = self.kg.query_triples(subject, predicate, object)
            return [
                {"subject": t.subject, "predicate": t.predicate, "object": t.object}
                for t in results
            ]

        # Visualization endpoints
        @self.app.post("/visualize")
        async def create_visualization(request: VisualizationRequest):
            """Create knowledge graph visualization"""
            try:
                # Create visualization
                filename = f"kg_visualization_{uuid.uuid4().hex[:8]}.html"
                filepath = os.path.join("visualization", filename)

                # Ensure directory exists
                os.makedirs("visualization", exist_ok=True)

                # Create network visualization
                net = Network(
                    height="600px", width="100%", bgcolor="#222222", font_color="white"
                )

                # Add entities
                entity_ids = list(self.kg.entities.keys())[: request.max_entities]
                for entity_id in entity_ids:
                    entity = self.kg.entities[entity_id]
                    net.add_node(
                        entity_id,
                        label=entity.name,
                        title=f"Type: {entity.entity_type}",
                        color=self._get_entity_color(entity.entity_type),
                    )

                # Add edges
                for triple in self.kg.triples:
                    if triple.subject in entity_ids and triple.object in entity_ids:
                        net.add_edge(
                            triple.subject,
                            triple.object,
                            label=triple.predicate,
                            color="#ffffff",
                        )

                # Save visualization
                net.save_graph(filepath)

                return {
                    "message": "Visualization created successfully",
                    "filepath": filepath,
                    "filename": filename,
                }

            except Exception as e:
                raise HTTPException(
                    status_code=500, detail=f"Visualization failed: {str(e)}"
                )

        @self.app.get("/visualize/{filename}")
        async def get_visualization(filename: str):
            """Get visualization file"""
            filepath = os.path.join("visualization", filename)
            if not os.path.exists(filepath):
                raise HTTPException(status_code=404, detail="Visualization not found")

            return FileResponse(filepath)

        # Embedding operations
        @self.app.post("/embeddings/train")
        async def train_embeddings(epochs: int = 50, lr: float = 0.001):
            """Train knowledge graph embeddings"""
            if not self.embedding_trainer:
                raise HTTPException(
                    status_code=503, detail="Embedding trainer not available"
                )

            try:
                losses = self.embedding_trainer.train(epochs=epochs, lr=lr)

                # Evaluate
                results = self.embedding_trainer.evaluate_link_prediction(k=3)

                return {
                    "message": "Embeddings trained successfully",
                    "final_loss": losses[-1] if losses else 0,
                    "hits_at_3": results.get("hits@3", 0),
                    "mean_rank": results.get("mean_rank", 0),
                }

            except Exception as e:
                raise HTTPException(
                    status_code=500, detail=f"Training failed: {str(e)}"
                )

        @self.app.get("/embeddings/similar/{entity_id}")
        async def get_similar_entities(entity_id: str, k: int = 5):
            """Get similar entities using embeddings"""
            if not self.embedding_trainer:
                raise HTTPException(
                    status_code=503, detail="Embedding trainer not available"
                )

            try:
                similar = self.embedding_trainer.find_similar_entities(entity_id, k=k)
                return {
                    "entity_id": entity_id,
                    "similar_entities": [
                        {
                            "entity_id": eid,
                            "name": self.kg.entities[eid].name,
                            "similarity": sim,
                        }
                        for eid, sim in similar
                    ],
                }
            except Exception as e:
                raise HTTPException(
                    status_code=500, detail=f"Similarity search failed: {str(e)}"
                )

        # Interactive web interface
        @self.app.get("/", response_class=HTMLResponse)
        async def web_interface():
            """Serve interactive web interface"""
            html_content = """
            <!DOCTYPE html>
            <html>
            <head>
                <title>Knowledge Graph API Interface</title>
                <style>
                    body { font-family: Arial, sans-serif; margin: 40px; }
                    .container { max-width: 800px; margin: 0 auto; }
                    .query-box { width: 100%; height: 100px; margin: 10px 0; }
                    .result-box { background: #f5f5f5; padding: 20px; margin: 10px 0; }
                    button { background: #007bff; color: white; padding: 10px 20px; border: none; cursor: pointer; }
                    button:hover { background: #0056b3; }
                </style>
            </head>
            <body>
                <div class="container">
                    <h1>🧠 Knowledge Graph API Interface</h1>
                    
                    <h2>Query Knowledge Graph</h2>
                    <textarea id="query" class="query-box" placeholder="Ask a question about the knowledge graph..."></textarea>
                    <br>
                    <button onclick="queryKG()">Query</button>
                    <button onclick="queryRAG()">Query with RAG</button>
                    
                    <div id="result" class="result-box" style="display: none;">
                        <h3>Result:</h3>
                        <div id="answer"></div>
                        <div id="reasoning"></div>
                    </div>
                    
                    <h2>Knowledge Graph Statistics</h2>
                    <div id="stats"></div>
                </div>
                
                <script>
                    async function queryKG() {
                        const query = document.getElementById('query').value;
                        if (!query) return;
                        
                        const response = await fetch('/query', {
                            method: 'POST',
                            headers: { 'Content-Type': 'application/json' },
                            body: JSON.stringify({ query: query })
                        });
                        
                        const result = await response.json();
                        displayResult(result);
                    }
                    
                    async function queryRAG() {
                        const query = document.getElementById('query').value;
                        if (!query) return;
                        
                        const response = await fetch('/rag/query', {
                            method: 'POST',
                            headers: { 'Content-Type': 'application/json' },
                            body: JSON.stringify({ query: query })
                        });
                        
                        const result = await response.json();
                        displayResult(result);
                    }
                    
                    function displayResult(result) {
                        document.getElementById('answer').innerHTML = '<strong>Answer:</strong> ' + result.answer;
                        document.getElementById('reasoning').innerHTML = '<strong>Reasoning:</strong><br>' + 
                            result.reasoning_steps.join('<br>');
                        document.getElementById('result').style.display = 'block';
                    }
                    
                    async function loadStats() {
                        const response = await fetch('/stats');
                        const stats = await response.json();
                        document.getElementById('stats').innerHTML = 
                            '<p>Entities: ' + stats.total_entities + '</p>' +
                            '<p>Relations: ' + stats.total_relations + '</p>' +
                            '<p>Triples: ' + stats.total_triples + '</p>';
                    }
                    
                    loadStats();
                </script>
            </body>
            </html>
            """
            return HTMLResponse(content=html_content)

    def _get_entity_color(self, entity_type: str) -> str:
        """Get color for entity type"""
        color_map = {
            "person": "#ff6b6b",
            "concept": "#4ecdc4",
            "field": "#45b7d1",
            "work": "#96ceb4",
            "discovery": "#feca57",
            "institution": "#ff9ff3",
            "country": "#a8e6cf",
        }
        return color_map.get(entity_type, "#d3d3d3")

    def run(self, host: str = "0.0.0.0", port: int = 8000):
        """Run the API server"""
        print(f"🚀 Starting Knowledge Graph API server on http://{host}:{port}")
        print(
            f"📊 Knowledge Graph stats: {len(self.kg.entities)} entities, {len(self.kg.triples)} triples"
        )
        print(f"🌐 Web interface: http://{host}:{port}")
        print(f"📚 API docs: http://{host}:{port}/docs")

        uvicorn.run(self.app, host=host, port=port)


def create_advanced_api(kg: KnowledgeGraph) -> AdvancedKnowledgeGraphAPI:
    """Create advanced API system for Knowledge Graph"""
    return AdvancedKnowledgeGraphAPI(kg)