Restore missing agent logic

agents/demystifier_agent.py

@@ -1,4 +1,3 @@
-# D:\jan-contract\agents\demystifier_agent.py
 
 import os
 from typing import TypedDict, List
@@ -11,13 +10,110 @@ from langchain_community.vectorstores import FAISS
 from langchain_core.prompts import PromptTemplate
 from langchain_core.runnables import RunnablePassthrough
 from langchain_core.output_parsers import StrOutputParser
-from core_utils.core_model_loaders import load_embedding_model
+from core_utils.core_model_loaders import load_embedding_model, load_gemini_llm
+from langgraph.graph import StateGraph, END
 
+# --- Initialize Models ---
 embedding_model = load_embedding_model()
+llm = load_gemini_llm()
 
-# ... (rest of imports)
+# --- 1. RAG Chain Logic ---
 
-# ...
+def create_rag_chain(retriever):
+    """Creates a RAG chain for answering questions about the document."""
+    
+    template = """Answer the question based only on the following context:
+    {context}
+
+    Question: {question}
+    """
+    prompt = PromptTemplate.from_template(template)
+
+    def format_docs(docs):
+        return "\n\n".join(doc.page_content for doc in docs)
+
+    rag_chain = (
+        {"context": retriever | format_docs, "question": RunnablePassthrough()}
+        | prompt
+        | llm
+        | StrOutputParser()
+    )
+    
+    return rag_chain
+
+# --- 2. Demystifier Graph Logic ---
+
+class DemystifierState(TypedDict):
+    document_chunks: List[str]
+    summary: str
+    key_terms: str
+    final_report: str
+
+def summarize_document(state: DemystifierState):
+    """Summarizes the provided document chunks."""
+    print("---NODE: Summarizing Document---")
+    text = "\n\n".join(state["document_chunks"])
+    # Truncate for safety if too large for prompt
+    text = text[:30000] 
+    
+    prompt = f"""
+    You are a legal expert. Summarize the following legal document content in simple, easy-to-understand language.
+    Focus on the main purpose and parties involved.
+    
+    Content:
+    {text}
+    """
+    response = llm.invoke(prompt)
+    return {"summary": response.content}
+
+def extract_key_terms(state: DemystifierState):
+    """Extracts and explains key legal terms."""
+    print("---NODE: Extracting Key Terms---")
+    text = "\n\n".join(state["document_chunks"])
+    text = text[:30000]
+    
+    prompt = f"""
+    Identify 5-7 complex legal terms or clauses from the text below.
+    List them and explain what they mean in plain English for a layperson.
+    
+    Content:
+    {text}
+    """
+    response = llm.invoke(prompt)
+    return {"key_terms": response.content}
+
+def generate_report(state: DemystifierState):
+    """Compiles the final analysis report."""
+    print("---NODE: Generating Final Report---")
+    
+    report = f"""
+    # Document Analysis
+    
+    ## 📝 Summary
+    {state['summary']}
+    
+    ## 🔑 Key Terms & Definitions
+    {state['key_terms']}
+    
+    ## 💡 Expert Advice
+    Always consult with a qualified lawyer for critical legal decisions. This analysis is AI-generated guidance.
+    """
+    return {"final_report": report}
+
+# --- Build the Graph ---
+workflow = StateGraph(DemystifierState)
+
+workflow.add_node("summarize", summarize_document)
+workflow.add_node("extract_terms", extract_key_terms)
+workflow.add_node("compile_report", generate_report)
+
+# Parallel execution of summary and terms
+workflow.set_entry_point("summarize")
+workflow.add_edge("summarize", "extract_terms")
+workflow.add_edge("extract_terms", "compile_report")
+workflow.add_edge("compile_report", END)
+
+demystifier_agent_graph = workflow.compile()
 
 # --- 4. The Master "Controller" Function ---
 def process_document_for_demystification(file_path: str):
@@ -41,7 +137,7 @@ def process_document_for_demystification(file_path: str):
     print("--- Running analysis graph for the report ---")
     chunk_contents = [chunk.page_content for chunk in chunks]
     # Limit context to avoid token limits if document is huge
-    graph_input = {"document_chunks": chunk_contents[:10]} 
+    graph_input = {"document_chunks": chunk_contents} 
     
     result = demystifier_agent_graph.invoke(graph_input)
     report = result.get("final_report")