ComputerAssignments/CA10_knowledge_graphs/src/reasoning/langgraph_reasoning.py

"""
Advanced Knowledge Graph Reasoning with LangGraph
This module implements sophisticated reasoning capabilities using LangGraph
"""

import os
import json
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
from datetime import datetime
import asyncio

# LangGraph and LangChain imports
from langgraph.graph import StateGraph, END
from langchain.schema import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import ChatPromptTemplate, PromptTemplate
from langchain.schema.runnable import RunnablePassthrough

# Vector search and embeddings
import chromadb
from sentence_transformers import SentenceTransformer
import numpy as np

# Graph processing
import networkx as nx
from pyvis.network import Network

# FastAPI for API endpoints
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn

import sys
import os

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from core.knowledge_graph import KnowledgeGraph, Triple
from core.scientific_kg import AdvancedScientificKG


@dataclass
class ReasoningState:
    """State for LangGraph reasoning workflow"""

    query: str
    context: List[str]
    entities: List[str]
    relations: List[str]
    reasoning_steps: List[str]
    final_answer: str
    confidence: float
    evidence: List[str]
    subgraph: Optional[KnowledgeGraph] = None


class KnowledgeGraphReasoner:
    """Advanced reasoning system using LangGraph"""

    def __init__(self, kg: KnowledgeGraph, model_name: str = "gemini-pro"):
        self.kg = kg
        self.model_name = model_name

        # Initialize LLM with Gemini API
        try:
            import os
            from dotenv import load_dotenv

            load_dotenv()

            google_api_key = os.getenv("GOOGLE_API_KEY")
            if google_api_key:
                self.llm = ChatGoogleGenerativeAI(
                    model=model_name, temperature=0.1, google_api_key=google_api_key
                )
                print(f"✅ Gemini LLM initialized: {model_name}")
            else:
                raise ValueError("GOOGLE_API_KEY not found in environment")
        except Exception as e:
            print(f"⚠️ LLM initialization failed: {e}")
            print("💡 Using mock LLM for demonstration")
            self.llm = None

        # Initialize embeddings and vector store
        self.embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2"
        )
        self.vector_store = None
        self._setup_vector_store()

        # Create reasoning workflow
        self.workflow = self._create_reasoning_workflow()

    def _setup_vector_store(self):
        """Setup vector store for semantic search"""
        # Convert KG triples to text for embedding
        documents = []
        metadatas = []

        for triple in self.kg.triples:
            doc_text = f"{triple.subject} {triple.predicate} {triple.object}"
            documents.append(doc_text)
            metadatas.append(
                {
                    "subject": triple.subject,
                    "predicate": triple.predicate,
                    "object": triple.object,
                    "triple_id": len(documents) - 1,
                }
            )

        # Create Chroma vector store
        self.vector_store = Chroma.from_texts(
            texts=documents, embedding=self.embeddings, metadatas=metadatas
        )

    def _create_reasoning_workflow(self) -> StateGraph:
        """Create LangGraph workflow for reasoning"""

        workflow = StateGraph(ReasoningState)

        # Add nodes
        workflow.add_node("parse_query", self._parse_query)
        workflow.add_node("extract_entities", self._extract_entities)
        workflow.add_node("find_relevant_context", self._find_relevant_context)
        workflow.add_node("reason_step_by_step", self._reason_step_by_step)
        workflow.add_node("generate_subgraph", self._generate_subgraph)
        workflow.add_node("synthesize_answer", self._synthesize_answer)

        # Add edges
        workflow.set_entry_point("parse_query")
        workflow.add_edge("parse_query", "extract_entities")
        workflow.add_edge("extract_entities", "find_relevant_context")
        workflow.add_edge("find_relevant_context", "reason_step_by_step")
        workflow.add_edge("reason_step_by_step", "generate_subgraph")
        workflow.add_edge("generate_subgraph", "synthesize_answer")
        workflow.add_edge("synthesize_answer", END)

        return workflow.compile()

    def _parse_query(self, state: ReasoningState) -> ReasoningState:
        """Parse and understand the query"""
        if self.llm is None:
            state.context.append(f"Query parsed (mock): {state.query}")
            return state

        prompt = ChatPromptTemplate.from_messages(
            [
                SystemMessage(
                    content="""You are a knowledge graph reasoning expert. 
            Parse the user query and identify the key components:
            1. What entities are mentioned?
            2. What relationships are being asked about?
            3. What type of reasoning is required?"""
                ),
                HumanMessage(content=f"Query: {state.query}"),
            ]
        )

        response = self.llm.invoke(prompt.format_messages())
        state.context.append(f"Query parsed: {response.content}")

        return state

    def _extract_entities(self, state: ReasoningState) -> ReasoningState:
        """Extract entities from the query"""
        if self.llm is None:
            # Fallback: simple extraction
            state.entities = [
                word
                for word in state.query.split()
                if word.lower() in [e.lower() for e in self.kg.entities.keys()]
            ]
            state.context.append(f"Extracted entities (mock): {state.entities}")
            return state

        prompt = ChatPromptTemplate.from_messages(
            [
                SystemMessage(
                    content="""Extract all entities mentioned in the query.
            Return them as a JSON list of strings."""
                ),
                HumanMessage(content=f"Query: {state.query}"),
            ]
        )

        response = self.llm.invoke(prompt.format_messages())

        try:
            entities = json.loads(response.content)
            state.entities = entities
        except:
            # Fallback: simple extraction
            state.entities = [
                word
                for word in state.query.split()
                if word.lower() in [e.lower() for e in self.kg.entities.keys()]
            ]

        state.context.append(f"Extracted entities: {state.entities}")
        return state

    def _find_relevant_context(self, state: ReasoningState) -> ReasoningState:
        """Find relevant context using vector search"""
        if not state.entities:
            return state

        # Search for relevant triples
        relevant_docs = self.vector_store.similarity_search(query=state.query, k=10)

        context_triples = []
        for doc in relevant_docs:
            metadata = doc.metadata
            context_triples.append(
                f"{metadata['subject']} {metadata['predicate']} {metadata['object']}"
            )

        state.context.extend(context_triples)
        state.context.append(f"Found {len(context_triples)} relevant triples")

        return state

    def _reason_step_by_step(self, state: ReasoningState) -> ReasoningState:
        """Perform step-by-step reasoning"""
        context_text = "\n".join(state.context)

        if self.llm is None:
            state.reasoning_steps.append(f"Mock reasoning for: {state.query}")
            return state

        prompt = ChatPromptTemplate.from_messages(
            [
                SystemMessage(
                    content="""You are a knowledge graph reasoning expert.
            Given the query and relevant context, perform step-by-step reasoning.
            Show your reasoning process clearly."""
                ),
                HumanMessage(
                    content=f"""
            Query: {state.query}
            
            Relevant Context:
            {context_text}
            
            Please reason step by step to answer the query.
            """
                ),
            ]
        )

        response = self.llm.invoke(prompt.format_messages())
        state.reasoning_steps.append(response.content)

        return state

    def _generate_subgraph(self, state: ReasoningState) -> ReasoningState:
        """Generate relevant subgraph"""
        if not state.entities:
            return state

        # Create subgraph from relevant entities
        subgraph = self.kg.get_subgraph(state.entities, max_hops=2)
        state.subgraph = subgraph

        state.context.append(
            f"Generated subgraph with {len(subgraph.entities)} entities"
        )

        return state

    def _synthesize_answer(self, state: ReasoningState) -> ReasoningState:
        """Synthesize final answer"""
        reasoning_text = "\n".join(state.reasoning_steps)
        context_text = "\n".join(state.context)

        if self.llm is None:
            state.final_answer = f"Mock answer for query: {state.query}"
            state.confidence = 0.5
            return state

        prompt = ChatPromptTemplate.from_messages(
            [
                SystemMessage(
                    content="""You are a knowledge graph reasoning expert.
            Based on the reasoning steps and context, provide a clear, 
            confident final answer to the query."""
                ),
                HumanMessage(
                    content=f"""
            Query: {state.query}
            
            Reasoning Steps:
            {reasoning_text}
            
            Context:
            {context_text}
            
            Provide a clear final answer with confidence score (0-1).
            """
                ),
            ]
        )

        response = self.llm.invoke(prompt.format_messages())
        state.final_answer = response.content
        state.confidence = 0.8  # Default confidence

        return state

    async def reason(self, query: str) -> Dict[str, Any]:
        """Main reasoning method"""
        initial_state = ReasoningState(
            query=query,
            context=[],
            entities=[],
            relations=[],
            reasoning_steps=[],
            final_answer="",
            confidence=0.0,
            evidence=[],
        )

        # Run the workflow
        final_state = await self.workflow.ainvoke(initial_state)

        return {
            "query": final_state.query,
            "answer": final_state.final_answer,
            "confidence": final_state.confidence,
            "reasoning_steps": final_state.reasoning_steps,
            "entities": final_state.entities,
            "evidence": final_state.evidence,
            "subgraph_size": (
                len(final_state.subgraph.entities) if final_state.subgraph else 0
            ),
        }


class MultiAgentReasoningSystem:
    """Multi-agent system for complex reasoning tasks"""

    def __init__(self, kg: KnowledgeGraph):
        self.kg = kg

        # Create specialized agents
        self.entity_agent = KnowledgeGraphReasoner(kg, "gemini-pro")
        self.relation_agent = KnowledgeGraphReasoner(kg, "gemini-pro")
        self.inference_agent = KnowledgeGraphReasoner(kg, "gemini-pro")

        # Create coordination workflow
        self.coordination_workflow = self._create_coordination_workflow()

    def _create_coordination_workflow(self) -> StateGraph:
        """Create workflow for multi-agent coordination"""

        workflow = StateGraph(ReasoningState)

        # Add agent nodes
        workflow.add_node("entity_analysis", self._entity_analysis)
        workflow.add_node("relation_analysis", self._relation_analysis)
        workflow.add_node("inference_synthesis", self._inference_synthesis)
        workflow.add_node("final_coordination", self._final_coordination)

        # Add edges
        workflow.set_entry_point("entity_analysis")
        workflow.add_edge("entity_analysis", "relation_analysis")
        workflow.add_edge("relation_analysis", "inference_synthesis")
        workflow.add_edge("inference_synthesis", "final_coordination")
        workflow.add_edge("final_coordination", END)

        return workflow.compile()

    async def _entity_analysis(self, state: ReasoningState) -> ReasoningState:
        """Entity-focused analysis"""
        result = await self.entity_agent.reason(state.query)
        state.context.extend(result["reasoning_steps"])
        state.entities = result["entities"]
        return state

    async def _relation_analysis(self, state: ReasoningState) -> ReasoningState:
        """Relation-focused analysis"""
        result = await self.relation_agent.reason(state.query)
        state.context.extend(result["reasoning_steps"])
        return state

    async def _inference_synthesis(self, state: ReasoningState) -> ReasoningState:
        """Inference synthesis"""
        result = await self.inference_agent.reason(state.query)
        state.context.extend(result["reasoning_steps"])
        return state

    async def _final_coordination(self, state: ReasoningState) -> ReasoningState:
        """Final coordination and answer synthesis"""
        # Synthesize all agent outputs
        state.final_answer = f"Multi-agent analysis complete. Query: {state.query}"
        state.confidence = 0.9
        return state

    async def reason(self, query: str) -> Dict[str, Any]:
        """Multi-agent reasoning"""
        initial_state = ReasoningState(
            query=query,
            context=[],
            entities=[],
            relations=[],
            reasoning_steps=[],
            final_answer="",
            confidence=0.0,
            evidence=[],
        )

        final_state = await self.coordination_workflow.ainvoke(initial_state)

        return {
            "query": final_state.query,
            "answer": final_state.final_answer,
            "confidence": final_state.confidence,
            "reasoning_steps": final_state.reasoning_steps,
            "entities": final_state.entities,
            "multi_agent": True,
        }


class KnowledgeGraphAPI:
    """FastAPI interface for Knowledge Graph reasoning"""

    def __init__(self, kg: KnowledgeGraph):
        self.kg = kg
        self.reasoner = KnowledgeGraphReasoner(kg)
        self.multi_agent = MultiAgentReasoningSystem(kg)

        # Create FastAPI app
        self.app = FastAPI(title="Knowledge Graph Reasoning API")
        self._setup_routes()

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

        @self.app.post("/reason")
        async def reason(query: str):
            """Single-agent reasoning"""
            try:
                result = await self.reasoner.reason(query)
                return result
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))

        @self.app.post("/reason-multi")
        async def reason_multi(query: str):
            """Multi-agent reasoning"""
            try:
                result = await self.multi_agent.reason(query)
                return result
            except Exception as e:
                raise HTTPException(status_code=500, detail=str(e))

        @self.app.get("/kg/stats")
        async def get_kg_stats():
            """Get knowledge graph statistics"""
            return {
                "entities": len(self.kg.entities),
                "relations": len(self.kg.relations),
                "triples": len(self.kg.triples),
            }

        @self.app.get("/kg/entities")
        async def get_entities():
            """Get all entities"""
            return list(self.kg.entities.keys())

        @self.app.post("/kg/query")
        async def query_kg(
            subject: str = None, predicate: str = None, object: str = None
        ):
            """Query knowledge graph"""
            results = self.kg.query_triples(subject, predicate, object)
            return [
                {"subject": t.subject, "predicate": t.predicate, "object": t.object}
                for t in results
            ]

    def run(self, host: str = "0.0.0.0", port: int = 8000):
        """Run the API server"""
        uvicorn.run(self.app, host=host, port=port)


def create_advanced_reasoning_system(
    kg: KnowledgeGraph,
) -> Tuple[KnowledgeGraphReasoner, MultiAgentReasoningSystem, KnowledgeGraphAPI]:
    """Create advanced reasoning system with LangGraph"""

    # Single agent reasoner
    reasoner = KnowledgeGraphReasoner(kg)

    # Multi-agent system
    multi_agent = MultiAgentReasoningSystem(kg)

    # API interface
    api = KnowledgeGraphAPI(kg)

    return reasoner, multi_agent, api