from typing import Dict, List, TypedDict, Annotated from langchain_core.messages import HumanMessage, AIMessage from langchain_core.prompts import ChatPromptTemplate from langchain_openai import ChatOpenAI, OpenAIEmbeddings from langchain_chroma import Chroma from langchain_core.output_parsers import StrOutputParser from langgraph.graph import Graph, StateGraph, END from ragas.metrics import faithfulness, answer_relevancy, context_precision, context_recall, answer_correctness from ragas import evaluate from datasets import Dataset import os from dotenv import load_dotenv import numpy as np import logging # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) load_dotenv() # Define the state class AgentState(TypedDict): messages: Annotated[List[HumanMessage | AIMessage], "The messages in the conversation"] context: Annotated[str, "The retrieved context"] response: Annotated[str, "The generated response"] metrics: Annotated[Dict, "The RAGAS metrics"] next: str # Initialize components llm = ChatOpenAI(model="gpt-3.5-turbo") embeddings = OpenAIEmbeddings(model="text-embedding-3-small") # Initialize Chroma with minimal configuration vectorstore = Chroma( persist_directory="./chroma_db", embedding_function=embeddings, collection_name="rag_collection", collection_metadata={"hnsw:space": "cosine"} ) # Define the retrieval function def retrieve(state: AgentState) -> Dict: try: messages = state["messages"] last_message = messages[-1] logger.info(f"Retrieving context for query: {last_message.content}") # Get relevant documents docs = vectorstore.similarity_search_with_score( last_message.content, k=10 # Increased number of documents ) if not docs: logger.warning("No relevant documents found in the knowledge base") raise ValueError("No relevant documents found in the knowledge base") # Filter and combine documents - using a lower threshold filtered_docs = [] for doc, score in docs: if score > 0.2: # Lower threshold for more context filtered_docs.append(doc) if not filtered_docs: logger.warning("No documents met the similarity threshold, using all retrieved documents") filtered_docs = [doc for doc, _ in docs] # Use all documents if none meet threshold # Sort documents by relevance (using the original scores) sorted_docs = sorted(zip(filtered_docs, [score for _, score in docs if score > 0.2]), key=lambda x: x[1], reverse=True) context = "\n\n".join([doc.page_content for doc, _ in sorted_docs]) logger.info(f"Using {len(filtered_docs)} documents for context") # Validate context if not context.strip(): logger.warning("No valid context could be retrieved") raise ValueError("No valid context could be retrieved") logger.info(f"Retrieved context length: {len(context)} characters") return { "context": context, "metrics": {}, # Initialize empty metrics "next": "generate" } except Exception as e: logger.error(f"Error in retrieval: {str(e)}") return { "context": "", "metrics": { "error": str(e), "faithfulness": 0.0, "answer_relevancy": 0.0, "context_precision": 0.0, "context_recall": 0.0, "answer_correctness": 0.0 }, "next": "generate" } # Define the generation function def generate(state: AgentState) -> Dict: try: messages = state["messages"] context = state["context"] if not context.strip(): logger.warning("Empty context in generation step") return { "response": "I apologize, but I couldn't find any relevant information in the knowledge base to answer your question. Please try rephrasing your question or upload more relevant documents.", "metrics": { "faithfulness": 0.0, "answer_relevancy": 0.0, "context_precision": 0.0, "context_recall": 0.0, "answer_correctness": 0.0, "note": "No context available for evaluation" }, "next": "evaluate" } logger.info("Generating response with context") # Create prompt with context prompt = ChatPromptTemplate.from_messages([ ("system", """You are a helpful assistant specialized in quantum computing. Use the following context to answer the question. Guidelines: 1. Base your answer strictly on the provided context 2. If the context doesn't contain relevant information, say so clearly 3. Be specific and technical in your explanations 4. Use bullet points or numbered lists when appropriate 5. Include relevant examples from the context 6. If discussing technical concepts, explain them clearly Context: {context}"""), ("human", "{question}") ]) chain = prompt | llm | StrOutputParser() response = chain.invoke({ "context": context, "question": messages[-1].content }) logger.info(f"Generated response length: {len(response)} characters") # Calculate metrics directly in generate try: logger.info("Creating dataset for metrics calculation") dataset = Dataset.from_dict({ "question": [messages[-1].content], "contexts": [[context]], "answer": [response], "ground_truth": [context] }) logger.info("Calculating RAGAS metrics") metrics_dict = {} result = evaluate(dataset, metrics=[faithfulness, answer_relevancy, context_precision, context_recall, answer_correctness]) metrics_dict["faithfulness"] = float(np.mean(result["faithfulness"])) metrics_dict["answer_relevancy"] = float(np.mean(result["answer_relevancy"])) metrics_dict["context_precision"] = float(np.mean(result["context_precision"])) metrics_dict["context_recall"] = float(np.mean(result["context_recall"])) metrics_dict["answer_correctness"] = float(np.mean(result["answer_correctness"])) logger.info(f"RAGAS metrics calculated: {metrics_dict}") except Exception as e: logger.error(f"Error calculating metrics: {str(e)}") metrics_dict = { "error": str(e), "faithfulness": 0.0, "answer_relevancy": 0.0, "context_precision": 0.0, "context_recall": 0.0, "answer_correctness": 0.0 } return { "response": response, "metrics": metrics_dict, "next": "evaluate" } except Exception as e: logger.error(f"Error in generation: {str(e)}") return { "response": "I apologize, but I encountered an error while generating a response. Please try again.", "metrics": { "error": str(e), "faithfulness": 0.0, "answer_relevancy": 0.0, "context_precision": 0.0, "context_recall": 0.0, "answer_correctness": 0.0 }, "next": "evaluate" } # Define the RAGAS evaluation function def evaluate_rag(state: AgentState) -> Dict: try: messages = state["messages"] context = state["context"] response = state["response"] # Detailed logging of input data logger.info("=== RAGAS Evaluation Debug Info ===") logger.info(f"Question: {messages[-1].content}") logger.info(f"Context length: {len(context)}") logger.info(f"Response length: {len(response)}") logger.info(f"Context preview: {context[:200]}...") logger.info(f"Response preview: {response[:200]}...") # Check if metrics are already in state if "metrics" in state: logger.info(f"Metrics found in state: {state['metrics']}") return {"context": context, "response": response, "metrics": state["metrics"], "next": END} # Validate inputs if not context.strip(): logger.error("Empty context detected") return {"context": context, "response": response, "metrics": { "faithfulness": 0.0, "answer_relevancy": 0.0, "context_precision": 0.0, "context_recall": 0.0, "answer_correctness": 0.0, "note": "Empty context" }, "next": END} if not response.strip(): logger.error("Empty response detected") return {"context": context, "response": response, "metrics": { "faithfulness": 0.0, "answer_relevancy": 0.0, "context_precision": 0.0, "context_recall": 0.0, "answer_correctness": 0.0, "note": "Empty response" }, "next": END} logger.info("Creating evaluation dataset...") try: # Create dataset for evaluation dataset = Dataset.from_dict({ "question": [messages[-1].content], "contexts": [[context]], "answer": [response], "ground_truth": [context] # Use context as ground truth for better evaluation }) logger.info("Dataset created successfully") # Initialize metrics dictionary metrics_dict = {} # Evaluate each metric separately try: result = evaluate(dataset, metrics=[faithfulness]) metrics_dict["faithfulness"] = float(np.mean(result["faithfulness"])) logger.info(f"Faithfulness calculated: {metrics_dict['faithfulness']}") except Exception as e: logger.error(f"Error calculating faithfulness: {str(e)}") metrics_dict["faithfulness"] = 0.0 try: result = evaluate(dataset, metrics=[answer_relevancy]) metrics_dict["answer_relevancy"] = float(np.mean(result["answer_relevancy"])) logger.info(f"Answer relevancy calculated: {metrics_dict['answer_relevancy']}") except Exception as e: logger.error(f"Error calculating answer_relevancy: {str(e)}") metrics_dict["answer_relevancy"] = 0.0 try: result = evaluate(dataset, metrics=[context_precision]) metrics_dict["context_precision"] = float(np.mean(result["context_precision"])) logger.info(f"Context precision calculated: {metrics_dict['context_precision']}") except Exception as e: logger.error(f"Error calculating context_precision: {str(e)}") metrics_dict["context_precision"] = 0.0 try: result = evaluate(dataset, metrics=[context_recall]) metrics_dict["context_recall"] = float(np.mean(result["context_recall"])) logger.info(f"Context recall calculated: {metrics_dict['context_recall']}") except Exception as e: logger.error(f"Error calculating context_recall: {str(e)}") metrics_dict["context_recall"] = 0.0 try: result = evaluate(dataset, metrics=[answer_correctness]) metrics_dict["answer_correctness"] = float(np.mean(result["answer_correctness"])) logger.info(f"Answer correctness calculated: {metrics_dict['answer_correctness']}") except Exception as e: logger.error(f"Error calculating answer_correctness: {str(e)}") metrics_dict["answer_correctness"] = 0.0 logger.info(f"RAGAS metrics calculated: {metrics_dict}") return {"context": context, "response": response, "metrics": metrics_dict, "next": END} except Exception as eval_error: logger.error(f"Error during RAGAS evaluation: {str(eval_error)}") logger.error(f"Error type: {type(eval_error)}") return {"context": context, "response": response, "metrics": { "faithfulness": 0.0, "answer_relevancy": 0.0, "context_precision": 0.0, "context_recall": 0.0, "answer_correctness": 0.0, "error": str(eval_error) }, "next": END} except Exception as e: logger.error(f"Error in RAGAS evaluation: {str(e)}") logger.error(f"Error type: {type(e)}") return {"context": context, "response": response, "metrics": { "faithfulness": 0.0, "answer_relevancy": 0.0, "context_precision": 0.0, "context_recall": 0.0, "answer_correctness": 0.0, "error": str(e) }, "next": END} # Create the workflow def create_rag_graph(): workflow = StateGraph(AgentState) # Add nodes workflow.add_node("retrieve", retrieve) workflow.add_node("generate", generate) workflow.add_node("evaluate", evaluate_rag) # Add edges workflow.add_edge("retrieve", "generate") workflow.add_edge("generate", "evaluate") workflow.add_edge("evaluate", END) # Set entry point workflow.set_entry_point("retrieve") # Compile app = workflow.compile() return app # Create the graph rag_graph = create_rag_graph()