src/agents/graph.py

import os
import sqlite3
import tempfile
import time
from typing import Annotated, TypedDict, Dict, Optional
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import FAISS
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage
from langchain_core.tools import tool
from langgraph.graph import START, END, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.checkpoint.sqlite import SqliteSaver

from src.agents.router import get_model
from src.memory.summariser import build_messages
from src.tools.code_executor import code_executor_tool

load_dotenv()

# Embeddings (shared)
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": "cpu"},
    encode_kwargs={"normalize_embeddings": True},
)

# Per-thread FAISS stores
_THREAD_RETRIEVERS: Dict[str, any] = {}
_THREAD_META: Dict[str, dict] = {}

# Web search tool
search_tool = DuckDuckGoSearchRun()

# PDF ingestion
def ingest_pdf(file_bytes: bytes, thread_id: str, filename: str):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as f:
        f.write(file_bytes)
        path = f.name

    loader = PyPDFLoader(path)
    docs = loader.load()
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = splitter.split_documents(docs)

    store = FAISS.from_documents(chunks, embeddings)
    _THREAD_RETRIEVERS[thread_id] = store.as_retriever(search_kwargs={"k": 4})
    _THREAD_META[thread_id] = {
        "filename": filename,
        "pages": len(docs),
        "chunks": len(chunks),
    }
    os.remove(path)
    return _THREAD_META[thread_id]

def get_thread_meta(thread_id: str):
    return _THREAD_META.get(thread_id)

# State
class State(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    query_type: str
    model_used: str
    latency_ms: float
SYSTEM_PROMPT = """You are a helpful research assistant.
When context is provided below, use it directly to answer the question.
Do not say "there is no document" if context is provided — just use it.
If no context is provided, answer from your own knowledge.
Never fabricate document sources.
Cite the source filename when answering from documents.
Be concise and accurate."""

# Agent node
def agent_node(state: State, config=None) -> dict:
    messages = state["messages"]
    thread_id = config["configurable"]["thread_id"] if config else "default"

    last_human = next(
        (m for m in reversed(messages) if isinstance(m, HumanMessage)), None
    )
    query = last_human.content if last_human else ""
    model_name, qtype = get_model(query)

    # Step 1:- manual RAG from per-thread FAISS
    context_block = ""
    source = ""
    retriever = _THREAD_RETRIEVERS.get(thread_id)
    if retriever:
        try:
            docs = retriever.invoke(query)
            if docs:
                context_block = "\n\n".join(d.page_content for d in docs)
                source = _THREAD_META.get(thread_id, {}).get("filename", "document")
        except Exception:
            pass
    # Step 2 — web search fallback if no doc context
    global search_tool
    if not context_block and search_tool is not None:
        try:
            web_result = search_tool.run(query)
            if web_result:
                context_block = web_result
                source = "web search"
        except Exception:
            pass

    # Step 3:- calculator for calc queries
    calc_result = ""
    if qtype == "calc":
        try:
            result = code_executor_tool.invoke({"code": f"print({query})"})
            if result.get("success"):
                calc_result = f"\nCalculation: {result['output']}"
        except Exception:
            pass

    # Build final prompt
    context_section = ""
    if context_block:
        context_section = f"\nSource: {source}\nContext:\n{context_block}\n"
    if calc_result:
        context_section += calc_result

    system = SystemMessage(content=SYSTEM_PROMPT + context_section)

    history = build_messages(messages[:-1], "")
    history = [m for m in history if not isinstance(m, SystemMessage)]
    final_messages = [system] + history + [HumanMessage(content=query)]

    llm = ChatGroq(
        model=model_name,
        api_key=os.getenv("GROQ_API_KEY"),
        temperature=0,
        max_tokens=1024,
    )

    t0 = time.perf_counter()
    response = llm.invoke(final_messages)
    latency_ms = (time.perf_counter() - t0) * 1000

    return {
        "messages": [response],
        "query_type": qtype,
        "model_used": model_name,
        "latency_ms": round(latency_ms, 1),
    }

# Graph
def build_graph():
    conn = sqlite3.connect("memory.db", check_same_thread=False)
    checkpointer = SqliteSaver(conn)

    graph = StateGraph(State)
    graph.add_node("agent", agent_node)
    graph.add_edge(START, "agent")
    graph.add_edge("agent", END)

    return graph.compile(checkpointer=checkpointer)

_graph = None

def get_graph():
    global _graph
    if _graph is None:
        _graph = build_graph()
    return _graph