added summarze agent

app/agents/context_agent.py

@@ -4,6 +4,7 @@ from langchain_groq import ChatGroq
 from app.prompts.context_agent_prompt import context_agent_template
 from app.tools.context_agent_tools import context_agent_tools
 from typing import Any
+from app.agent_memory_store import memory_store
 
 context_agent = create_agent(
     model=ChatGroq(
@@ -11,7 +12,7 @@ context_agent = create_agent(
         temperature=0.1,
     ),
     tools=context_agent_tools,
-    store=p_store,  
+    store=memory_store,  
     middleware=[
         ToolCallLimitMiddleware[Any,None](
            tool_name="search_memory",

app/agents/summarizer_agent.py

@@ -0,0 +1,8 @@
+from langchain_groq import ChatGroq
+
+summarizer_agent=ChatGroq(
+
+    model="llama-3.3-70b-versatile",
+    temperature=0.1
+    
+)

app/graph.py

@@ -13,8 +13,9 @@ from app.nodes.store_memory_data_node import store_memory_and_data_node
 from app.nodes.unsafe_email_node import unsafe_emails_node
 from app.nodes.check_email_exist_node import *
 from langgraph.types import RetryPolicy
+from app.nodes.summarise_email_body_node import summarise_email_body_node
+from app.nodes.check_token_count_node import *
 from psycopg import OperationalError # Or sqlalchemy.exc.OperationalError depending on your driver
-# imoprt display
 from IPython.display import Image, display
 
 
@@ -38,11 +39,13 @@ builder = StateGraph(EmailAgentState)
 # Nodes
 builder.add_node("safety_check_node", safety_classifier_node)  
 builder.add_node("check_previous_email_exist_node", check_previous_email_exist_node)
+builder.add_node("check_token_count_node", check_token_count_node)
+builder.add_node("summarise_email_body_node", summarise_email_body_node)
 builder.add_node("triage_node", triage_node)
 builder.add_node("prepare_context_node", prepare_context_node)
 builder.add_node("email_writing_agent", email_writing_agent_node)
 
-
+# --- APPLY RETRY POLICIES HERE ---
 builder.add_node(
     "store_memory_and_data_node", 
     store_memory_and_data_node, 
@@ -61,10 +64,27 @@ builder.add_node("tools", ToolNode(email_writing_agent_tools), retry_policy=tool
 # Edges (Same as your original logic)
 builder.add_edge(START, "safety_check_node")
 
-builder.add_conditional_edges("safety_check_node", after_safety, {
-    "unsafe_emails_node": "unsafe_emails_node",
-    "triage_node": "triage_node",
-})
+builder.add_conditional_edges(
+    "safety_check_node", 
+    after_safety, 
+    {
+        "unsafe": "unsafe_emails_node",
+        "safe": "check_token_count_node" 
+    }
+)
+
+builder.add_conditional_edges(
+    "check_token_count_node", 
+    check_token_limit_router, 
+    {
+        "summarize": "summarise_email_body_node",
+        "triage": "triage_node" 
+    }
+)
+
+builder.add_edge("summarise_email_body_node", "triage_node")
+
+
 
 builder.add_conditional_edges("triage_node", route_after_triage, {
     "check_previous_email_exist_node": "check_previous_email_exist_node", 
@@ -104,9 +124,3 @@ builder.add_edge("parse_node", "store_memory_and_data_node")
 builder.add_edge("store_memory_and_data_node", END)
 builder.add_edge("unsafe_emails_node", END)
 builder.add_edge("archive_node", END)
-
-
-
-graph = builder.compile() 
-
-display(graph)

app/nodes/check_token_count_node.py

@@ -0,0 +1,36 @@
+from app.state.state import EmailAgentState
+from langchain_groq import ChatGroq
+
+llm_for_token_count=ChatGroq(
+    model="meta-llama/llama-4-scout-17b-16e-instruct",
+    temperature=0.1,
+)
+
+def count_input_tokens(subject:str ,body:str) -> int:
+    text=subject+body
+    return llm_for_token_count.get_num_tokens(text)
+
+def check_token_count_node(state: EmailAgentState):
+    """
+    This is a formal Node. It calculates tokens and can 
+    update the state if you add a 'token_count' key to your TypedDict.
+    """
+    subject = state.get('sender_subject', "")
+    body = state.get('sender_email_body', "") 
+    
+    tokens = count_input_tokens(subject, body)
+    print(f"--- NODE: Token Count calculated as {tokens} ---")
+    
+    # We return the count so it's stored in the state for the next router to see
+    return {"sender_email_token_count": tokens}
+
+
+def check_token_limit_router(state: EmailAgentState):
+    """
+    Acts as a router to decide if we need summarization.
+    """
+    if state.get("sender_email_token_count", 0) > 100000:
+
+        return "summarize"
+    
+    return "triage"

app/nodes/safety_check_node.py

@@ -34,8 +34,10 @@ def safety_classifier_node(state: EmailAgentState) -> dict:
         return {"is_safe": False, "safety_reason": f"API error {e.response.status_code}"}
     
 
-
 def after_safety(state: EmailAgentState) -> str:
-    if not state["is_safe"]:
-        return "unsafe_emails_node"
-    return "triage_node"
+    """
+    Only handles safety routing.
+    """
+    if not state.get("is_safe"):
+        return "unsafe"
+    return "safe"

app/nodes/summarise_email_body_node.py

@@ -0,0 +1,36 @@
+from app.state.state import EmailAgentState
+from langchain_core.prompts import PromptTemplate
+from langchain_classic.chains.summarize import load_summarize_chain
+from langchain_classic.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+from langchain_classic.prompts import PromptTemplate
+from app.agents.summarizer_agent import summarizer_agent
+from app.prompts.summarizer_agent_prompt import summarize_agent_initial_prompt_template,summarize_agent_refine_template
+def summarise_email_body(body: str):
+    # Tip: 100 is very small for an email; 500-1000 is usually better
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    docs = [Document(page_content=t) for t in text_splitter.split_text(body)]
+
+    chain = load_summarize_chain(
+        llm=summarizer_agent, 
+        chain_type="refine",
+        question_prompt=summarize_agent_initial_prompt,
+        refine_prompt=summarize_agent_refine_prompt,
+        document_variable_name="body"  # <--- ADD THIS LINE
+    )
+    summary = chain.invoke(docs)
+    return summary['output_text']
+
+
+def summarise_email_body_node(state:EmailAgentState)->dict:
+
+
+    subject=state['sender_subject']
+    body=state['sender_email_body']
+
+
+
+    summary=summarise_email_body(body)
+
+    return {"sender_email_body":summary}
+

app/prompts/summarizer_agent_prompt.py

@@ -0,0 +1,25 @@
+
+
+from langchain_core.prompts import PromptTemplate
+
+summarize_agent_initial_prompt_template = """Write a summary of the following email body. 
+Include links (if any) and key info (req for reply) if given and maintain a email  structure showing who sent it:
+
+"{body}"
+
+CONCISE SUMMARY:"""
+
+
+summarize_agent_refine_template = """
+We have an existing summary: {existing_answer}
+We have more email content below:
+------------
+{body}
+------------
+Given the new context, refine the summary. 
+Ensure you keep the repository links and the sender's info.
+"""
+
+summarize_agent_initial_prompt = PromptTemplate(template=summarize_agent_initial_prompt_template, input_variables=["body"])
+
+summarize_agent_refine_prompt = PromptTemplate(template=summarize_agent_refine_template, input_variables=["existing_answer", "body"])

app/state/state.py

@@ -21,9 +21,12 @@ class EmailAgentState(TypedDict):
 
     sender_email_id:         str
 
-    sender_subject:        str     
+    sender_subject:        str   
+
     
-    user_name:      str                  
+    user_name:      str      
+
+    sender_email_token_count: Optional[int]            
 
     # Safety node output
     is_safe:        Optional[bool]