app.py

from fastapi import FastAPI, HTTPException
import os
from typing import List, Dict
from dotenv import load_dotenv
import logging
from pathlib import Path

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Qdrant as QdrantVectorStore
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_groq import ChatGroq
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from qdrant_client.models import PointIdsList

from langgraph.graph import MessagesState, StateGraph
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import END
from langgraph.prebuilt import tools_condition
from langgraph.checkpoint.memory import MemorySaver

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

load_dotenv()
GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
GROQ_API_KEY = os.getenv('GROQ_API_KEY')

if not GOOGLE_API_KEY or not GROQ_API_KEY:
    raise ValueError("API keys not set in environment variables")

app = FastAPI()

class QASystem:
    def __init__(self):
        self.vector_store = None
        self.graph = None
        self.memory = None
        self.embeddings = None
        self.client = None
        self.pdf_dir = "pdfss"

    def load_pdf_documents(self):
        documents = []
        pdf_dir = Path(self.pdf_dir)
        
        if not pdf_dir.exists():
            raise FileNotFoundError(f"PDF directory not found: {self.pdf_dir}")
        
        for pdf_path in pdf_dir.glob("*.pdf"):
            try:
                loader = PyPDFLoader(str(pdf_path))
                documents.extend(loader.load())
                logger.info(f"Loaded PDF: {pdf_path}")
            except Exception as e:
                logger.error(f"Error loading PDF {pdf_path}: {str(e)}")

        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=100
        )
        split_docs = text_splitter.split_documents(documents)
        logger.info(f"Split documents into {len(split_docs)} chunks")
        return split_docs

    def initialize_system(self):
        try:
            self.client = QdrantClient(":memory:")
            
            try:
                self.client.get_collection("pdf_data")
            except Exception:
                self.client.create_collection(
                    collection_name="pdf_data",
                    vectors_config=VectorParams(size=768, distance=Distance.COSINE),
                )
                logger.info("Created new collection: pdf_data")
            
            self.embeddings = GoogleGenerativeAIEmbeddings(
                model="models/embedding-001",
                google_api_key=GOOGLE_API_KEY
            )
            
            self.vector_store = QdrantVectorStore(
                client=self.client,
                collection_name="pdf_data",
                embeddings=self.embeddings,
            )
            
            documents = self.load_pdf_documents()
            if documents:
                try:
                    points = self.client.scroll(collection_name="pdf_data", limit=100)[0]
                    if points:
                        self.client.delete(
                            collection_name="pdf_data",
                            points_selector=PointIdsList(
                                points=[p.id for p in points]
                            )
                        )
                except Exception as e:
                    logger.error(f"Error clearing vectors: {str(e)}")
                
                self.vector_store.add_documents(documents)
                logger.info(f"Added {len(documents)} documents to vector store")

            llm = ChatGroq(
                model="llama3-8b-8192", 
                api_key=GROQ_API_KEY,
                temperature=0.7
            )
            
            graph_builder = StateGraph(MessagesState)

            def query_or_respond(state: MessagesState):
                retrieved_docs = [m for m in state["messages"] if m.type == "tool"]
    
                if retrieved_docs:
                    context = ' '.join(m.content for m in retrieved_docs)
                else:
                    context = "mountain bicycle documentation knowledge"

                system_prompt = (
                    "You are an AI assistant embedded within the Interactive Electronic Technical Manual (IETM) for Mountain Cycles.. "
                    "Always provide accurate responses with references to provided data. "
                    "If the user query is not technical-specific, still respond from a IETM perspective."
                    f"\n\nContext:\n{context}"
                )
    
                messages = [SystemMessage(content=system_prompt)] + state["messages"]
    
                logger.info(f"Sending to LLM: {[m.content for m in messages]}")  # Debugging log

                response = llm.invoke(messages)
                return {"messages": [response]}

            def generate(state: MessagesState):
                retrieved_docs = [m for m in reversed(state["messages"]) if m.type == "tool"][::-1]
    
                context = ' '.join(m.content for m in retrieved_docs) if retrieved_docs else "mountain bicycle documentation knowledge"

                system_prompt = (
                    "You are an AI assistant embedded within the Interactive Electronic Technical Manual (IETM) for Mountain Cycles. "
                    "Your responses MUST be accurate, concise (5 sentences max)."
                    f"\n\nContext:\n{context}"
                )

                messages = [SystemMessage(content=system_prompt)] + state["messages"]
    
                logger.info(f"Sending to LLM: {[m.content for m in messages]}")  # Debugging log

                response = llm.invoke(messages)
                return {"messages": [response]}


            graph_builder.add_node("query_or_respond", query_or_respond)
            graph_builder.add_node("generate", generate)
            
            graph_builder.set_entry_point("query_or_respond")
            graph_builder.add_edge("query_or_respond", "generate")
            graph_builder.add_edge("generate", END)
            
            self.memory = MemorySaver()
            self.graph = graph_builder.compile(checkpointer=self.memory)
            return True

        except Exception as e:
            logger.error(f"System initialization error: {str(e)}")
            return False

    def process_query(self, query: str) -> List[Dict[str, str]]:
        try:
            responses = []
            for step in self.graph.stream(
                {"messages": [HumanMessage(content=query)]},
                stream_mode="values",
                config={"configurable": {"thread_id": "abc123"}}
            ):
                if step["messages"]:
                    responses.append({
                        'content': step["messages"][-1].content,
                        'type': step["messages"][-1].type
                    })
            return responses
        except Exception as e:
            logger.error(f"Query processing error: {str(e)}")
            return [{'content': f"Query processing error: {str(e)}", 'type': 'error'}]

qa_system = QASystem()
if qa_system.initialize_system():
    logger.info("QA System Initialized Successfully")
else:
    raise RuntimeError("Failed to initialize QA System")
@app.post("/query")
async def query_api(query: str):
    responses = qa_system.process_query(query)
    return {"responses": responses}