Spaces:
Sleeping
Sleeping
| import datetime | |
| import os | |
| import re | |
| from typing import TypedDict, Optional, List | |
| from llama_index.core.schema import Document | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| from langchain_openai import ChatOpenAI | |
| from llama_index.core.retrievers import BaseRetriever | |
| from langgraph.graph import StateGraph, END | |
| class GraphState(TypedDict): | |
| query: str | |
| response: Optional[str] | |
| retrieved_docs: Optional[List[Document]] | |
| context: Optional[str] | |
| history: list | |
| retriever: Optional[BaseRetriever] | |
| llm: any | |
| def _select_chunks_by_relevance_and_position(docs: List[Document], query: str, target_count: int) -> List[Document]: | |
| """Select chunks based on semantic relevance + positional importance without summarization.""" | |
| if len(docs) <= target_count: | |
| return docs | |
| selected = [] | |
| if len(docs) >= 2: | |
| selected = [docs[0], docs[-1]] | |
| remaining_slots = target_count - 2 | |
| middle_docs = docs[1:-1] | |
| else: | |
| remaining_slots = target_count | |
| middle_docs = docs[:] | |
| if remaining_slots > 0 and middle_docs: | |
| query_words = set(query.lower().split()) | |
| scored_docs = [] | |
| for doc in middle_docs: | |
| content = doc.text.lower() | |
| relevance_score = sum(content.count(word) for word in query_words) | |
| length_bonus = len(doc.text) / 1000 | |
| total_score = relevance_score + length_bonus | |
| scored_docs.append((total_score, doc)) | |
| scored_docs.sort(key=lambda x: x[0], reverse=True) | |
| selected.extend([doc for _, doc in scored_docs[:remaining_slots]]) | |
| return selected[:target_count] | |
| def _create_context_with_priorities(docs: List[Document], max_chars: int) -> str: | |
| """Create context by prioritizing important chunks without summarization.""" | |
| if not docs: | |
| return "" | |
| contexts = [] | |
| total_chars = 0 | |
| priority_docs = [] | |
| if len(docs) >= 2: | |
| priority_docs.append(("BOUNDARY", docs[0])) | |
| priority_docs.append(("BOUNDARY", docs[-1])) | |
| middle_docs = docs[1:-1] if len(docs) > 2 else [] | |
| middle_docs.sort(key=lambda d: len(d.text), reverse=True) | |
| priority_docs.extend([("CONTENT", doc) for doc in middle_docs]) | |
| else: | |
| priority_docs = [("CONTENT", doc) for doc in docs] | |
| for priority_type, doc in priority_docs: | |
| content = doc.text | |
| if total_chars + len(content) > max_chars: | |
| remaining_chars = max_chars - total_chars | |
| if remaining_chars > 200: | |
| truncated = content[:remaining_chars] | |
| last_period = truncated.rfind('.') | |
| if last_period > remaining_chars * 0.8: | |
| truncated = truncated[:last_period + 1] | |
| truncated += "...[truncated]" | |
| contexts.append(f"[{priority_type}] {truncated}") | |
| break | |
| contexts.append(f"[{priority_type}] {content}") | |
| total_chars += len(content) | |
| return "\n\n---\n\n".join(contexts) | |
| def router_node(state: GraphState) -> GraphState: | |
| """ | |
| Router that determines the next step based on available retriever. | |
| """ | |
| print("---NODE: ROUTER---") | |
| return state | |
| def general_chat_node(state: GraphState) -> GraphState: | |
| """ | |
| Generates a response for general conversation using the LLM. | |
| """ | |
| print("---NODE: GENERAL CHAT---") | |
| llm = state["llm"] | |
| query = state["query"] | |
| history = state["history"] | |
| history_context = "\n".join([f"Human: {q}\nAI: {a}" for q, a in history]) | |
| current_time = datetime.datetime.now().strftime("%Y-%m-%d %I:%M %p") | |
| prompt = f"""You are Sarathi, a friendly and knowledgeable AI assistant. | |
| Instructions: | |
| - The current time is {current_time}. | |
| - Be conversational, helpful, and concise. | |
| - If you don't know the answer to a question, admit it honestly. | |
| - Do not make up information. | |
| **Conversation History:** | |
| {history_context} | |
| **Current User Query:** | |
| Human: {query} | |
| AI:""" | |
| try: | |
| response_obj = llm.invoke(prompt) | |
| response_text = getattr(response_obj, "content", str(response_obj)) | |
| except Exception as e: | |
| response_text = f"Model inference failed: {str(e)}" | |
| return {"response": response_text.strip()} | |
| def retrieve_node(state: GraphState) -> GraphState: | |
| """ | |
| Retrieves relevant documents and creates intelligent context without summarization. | |
| """ | |
| print("---NODE: RETRIEVE---") | |
| query = state["query"] | |
| retriever = state["retriever"] | |
| history = state["history"] | |
| context = "" | |
| retrieved_docs = [] | |
| try: | |
| q_len = len(query.split()) | |
| if q_len < 5: | |
| top_k = 5 | |
| elif q_len < 15: | |
| top_k = 8 | |
| else: | |
| top_k = 12 | |
| retrieved_docs = retriever.retrieve(query) | |
| print(f"Retrieved {len(retrieved_docs)} documents") | |
| if retrieved_docs: | |
| max_chunks = min(len(retrieved_docs), top_k) | |
| selected_docs = _select_chunks_by_relevance_and_position( | |
| retrieved_docs, query, max_chunks | |
| ) | |
| print(f"Selected {len(selected_docs)} chunks for context") | |
| MAX_CONTEXT_CHARS = 6000 | |
| doc_context = _create_context_with_priorities(selected_docs, MAX_CONTEXT_CHARS) | |
| if history and len(doc_context) < MAX_CONTEXT_CHARS * 0.8: | |
| history_context = "\n\n".join([f"Human: {q}\nAI: {a}" for q, a in history[-3:]]) # Last 3 exchanges | |
| remaining_chars = MAX_CONTEXT_CHARS - len(doc_context) | |
| if len(history_context) <= remaining_chars: | |
| context = f"{doc_context}\n\n--- Recent Chat History ---\n{history_context}" | |
| else: | |
| context = doc_context | |
| else: | |
| context = doc_context | |
| else: | |
| context = "No relevant content found in the document." | |
| except Exception as e: | |
| print(f"Error in retrieve_node: {e}") | |
| context = f"Retriever failed: {str(e)}" | |
| retrieved_docs = [] | |
| return {"retrieved_docs": retrieved_docs, "context": context} | |
| def generate_rag_node(state: GraphState) -> GraphState: | |
| """ | |
| Generates an answer using the retrieved context from the PDF. | |
| """ | |
| print("---NODE: GENERATE RAG RESPONSE---") | |
| llm = state["llm"] | |
| query = state["query"] | |
| context = state.get("context", "") | |
| if not context.strip(): | |
| return {"response": "I don't have enough context from the document to answer this question."} | |
| prompt = f"""You are Sarathi, an intelligent document analysis assistant. Use the provided context to answer the user's question thoroughly and accurately. | |
| Context from document: | |
| {context} | |
| Question: {query} | |
| Instructions: | |
| - Provide detailed, well-structured answers | |
| - If the question has multiple parts, address each part separately | |
| - Cite specific information from the context when possible | |
| - If the context doesn't fully cover the question, clearly state what's missing | |
| - Use a conversational but professional tone | |
| Answer:""" | |
| try: | |
| response_obj = llm.invoke(prompt) | |
| response_text = getattr(response_obj, "content", str(response_obj)) | |
| except Exception as e: | |
| response_text = f"Model inference failed: {str(e)}" | |
| return {"response": response_text.strip()} | |
| def route_query(state: GraphState) -> str: | |
| """ | |
| Checks if a retriever is available in the state to decide the next step. | |
| """ | |
| if state.get("retriever"): | |
| print("---ROUTING: PDF available, routing to retrieve_node---") | |
| return "retrieve" | |
| else: | |
| print("---ROUTING: No PDF, routing to general_chat_node---") | |
| return "general_chat" | |
| def build_graph(model_type: str = "groq", retriever=None): | |
| """ | |
| Builds the workflow graph with LLM and retriever. | |
| """ | |
| def make_llm(temp: float): | |
| if model_type == "groq": | |
| api_key = os.getenv("OPENROUTER_API_KEY", "").strip() | |
| if not api_key: | |
| raise ValueError("OPENROUTER_API_KEY environment variable not set.") | |
| return ChatOpenAI( | |
| model="x-ai/grok-4-fast", | |
| base_url="https://openrouter.ai/api/v1", | |
| api_key=api_key, | |
| temperature=temp, | |
| ) | |
| elif model_type == "gemini": | |
| api_key = os.getenv("GEMINI_API_KEY", "").strip() | |
| if not api_key: | |
| raise ValueError("GEMINI_API_KEY environment variable not set.") | |
| return ChatGoogleGenerativeAI( | |
| model="gemini-2.0-flash", | |
| api_key=api_key, | |
| temperature=temp, | |
| ) | |
| else: | |
| raise ValueError("Invalid model_type. Choose 'groq' or 'gemini'.") | |
| workflow = StateGraph(GraphState) | |
| workflow.add_node("router", router_node) | |
| workflow.add_node("general_chat", general_chat_node) | |
| workflow.add_node("retrieve", retrieve_node) | |
| workflow.add_node("generate", generate_rag_node) | |
| workflow.set_entry_point("router") | |
| workflow.add_conditional_edges( | |
| "router", | |
| route_query, | |
| { | |
| "retrieve": "retrieve", | |
| "general_chat": "general_chat", | |
| }, | |
| ) | |
| workflow.add_edge("retrieve", "generate") | |
| workflow.add_edge("generate", END) | |
| workflow.add_edge("general_chat", END) | |
| compiled_graph = workflow.compile() | |
| def graph_wrapper(query: str, temperature: float = 0.7, raw_text: str = None, history=None, retriever_override=None): | |
| active_retriever = retriever_override or retriever | |
| llm = make_llm(temperature) | |
| return compiled_graph.invoke({ | |
| "query": query, | |
| "retriever": active_retriever.as_retriever() if active_retriever else None, | |
| "history": history or [], | |
| "llm": llm, | |
| "response": None, | |
| "retrieved_docs": None, | |
| "context": None, | |
| }) | |
| return graph_wrapper |