import tiktoken from langchain_core.messages import trim_messages,HumanMessage, AIMessage from langgraph.graph import MessagesState,StateGraph, START, END, MessagesState from langgraph.prebuilt import ToolNode, tools_condition from pydantic import BaseModel, Field from typing import Literal import json import logging from back_end.llm import LLMConfig, build_chat_model, strip_thinking_blocks, with_structured_output logger = logging.getLogger(__name__) from back_end.config import ( SUPERVISOR_SYSTEM_PROMPT, AGENT_SYSTEM_PROMPT_HEADER, AGENT_SYSTEM_PROMPT_TOOLS, AGENT_SYSTEM_PROMPT_TOOLS_KEYWORD_ONLY, AGENT_SYSTEM_PROMPT_TOOLS_NO_DB, AGENT_SYSTEM_PROMPT_FOOTER, ) enc = tiktoken.get_encoding("cl100k_base") def _tiktoken_counter(messages): total = 0 for m in messages: text_to_encode = "" # 1. Extract content and tool_calls safely if isinstance(m, dict): content = m.get("content", "") tool_calls = m.get("tool_calls", []) else: content = getattr(m, "content", "") tool_calls = getattr(m, "tool_calls", []) # 2. Handle string or list content if isinstance(content, list): text_to_encode += str(content) else: text_to_encode += str(content) # 3. CRITICAL: Catch tool calls so they don't bypass the counter if tool_calls: text_to_encode += json.dumps(tool_calls) # Encode and count total += len(enc.encode(text_to_encode)) return total # --------------------------------------------------------- # 1. AGENT NODE # --------------------------------------------------------- def initialize_agent(tools: list, llm_cfg: LLMConfig): llm = build_chat_model(llm_cfg, temperature=0) # Describe exactly the tools this repo has (semantic search is skipped for big repos) tool_names = {t.name for t in tools} if "semantic_code_search" in tool_names: tools_prompt = AGENT_SYSTEM_PROMPT_TOOLS elif "keyword_code_search" in tool_names: tools_prompt = AGENT_SYSTEM_PROMPT_TOOLS_KEYWORD_ONLY else: tools_prompt = AGENT_SYSTEM_PROMPT_TOOLS_NO_DB system_prompt = f"{AGENT_SYSTEM_PROMPT_HEADER}\n\n{tools_prompt}\n\n{AGENT_SYSTEM_PROMPT_FOOTER}" llm_with_tools = llm.bind_tools(tools) message_trimmer = trim_messages( max_tokens=200000, strategy="last", token_counter=_tiktoken_counter, # We Use the Gemini model's specific token counter but it will make http request which will take too long so just just tiktoken wich will be good enough include_system=True, # NEVER delete the system prompt/repo map allow_partial=False # Don't chop a message in half ) # Call the model to generate a response based on the current state. # Given the question, it will decide to retrieve using the retriever tool, or simply respond to the user. def generate_query_or_respond(state: MessagesState): # 1. Inject the system prompt into the message history messages_to_evaluate = [{"role": "system", "content": system_prompt}] + state["messages"] # 2. to save context window,or not to runout of tokens we trim the context from past which in above max limit that we trimmed_messages = message_trimmer.invoke(messages_to_evaluate) # 3. Generate the response (PASS IN THE TRIMMED MESSAGES) response = llm_with_tools.invoke(trimmed_messages) return {"messages": [response]} return generate_query_or_respond # --------------------------------------------------------- # 2. THE LEAD ARCHITECT (SUPERVISOR NODE) # --------------------------------------------------------- # 1. Define the decision schema class SupervisorDecision(BaseModel): reasoning: str = Field( description="1. What did the user ask? 2. What raw data is in the tool outputs? 3. Is the raw data sufficient to answer the user?" ) status: Literal["ACCEPT", "REJECT"] = Field( description="ACCEPT if the RAW TOOL OUTPUTS contain enough info to answer the user. REJECT if the agent needs to search for more specific files." ) content: str = Field( description="If ACCEPT: Write the final, exhaustive response to the user. If REJECT: Write targeted instructions telling the agent what to search for next." ) def initialize_supervisor(llm_cfg: LLMConfig): powerful_llm = build_chat_model(llm_cfg, temperature=0.2, max_tokens=16000) powerful_agent = with_structured_output(llm_cfg, powerful_llm, SupervisorDecision) def supervisor_node(state: MessagesState): # Calculate iteration count based on previous feedback messages iteration_count = sum( 1 for m in state["messages"] if isinstance(m, HumanMessage) and "SUPERVISOR FEEDBACK:" in m.content ) system_prompt = SUPERVISOR_SYSTEM_PROMPT # STRUCTURAL SAFEGUARD: Force accept after 2 rejections if iteration_count >= 2: system_prompt += """ \n\n*** CRITICAL OVERRIDE *** You have rejected the researcher 2 times. You MUST now output status="ACCEPT" and synthesize the best possible final answer from ALL available evidence, explicitly noting what is implicit vs explicit. DO NOT REJECT. """ # The researcher's reasoning blocks were produced under a different system # prompt, so they can't be replayed here (Claude rejects them); drop them. messages_to_evaluate = [ {"role": "system", "content": system_prompt} ] + strip_thinking_blocks(state["messages"]) try: decision = powerful_agent.invoke(messages_to_evaluate) except Exception as e: # Don't echo raw provider errors to users; main.py maps them to a safe message logger.warning("Supervisor error: %s", type(e).__name__) raise if decision.status == "ACCEPT": return { "messages": [ AIMessage(content=decision.content) ] } return { "messages": [ HumanMessage( content=f"SUPERVISOR FEEDBACK: {decision.content}" ) ] } return supervisor_node # --- Custom Router for the Supervisor --- def route_supervisor(state: MessagesState): last_message = state["messages"][-1] # If the supervisor returned an AIMessage, it ACCEPTED the work. We are done. if isinstance(last_message, AIMessage): return END # If it returned a HumanMessage, it REJECTED the work. Send back to the researcher. return "agent" def build_workflow(tools: list, llm_cfg: LLMConfig): """Compile the agent graph for one chat request using the caller's own LLM key.""" agent_node = initialize_agent(tools, llm_cfg) supervisor_node = initialize_supervisor(llm_cfg) # --- Building the Graph --- workflow = StateGraph(MessagesState) # --- Add our nodes to the graph --- # Set the entry point: Start by calling the agent workflow.add_edge(START, "agent") workflow.add_node("agent", agent_node) workflow.add_node("tools", ToolNode(tools)) workflow.add_node("supervisor",supervisor_node) # --- Routing --- # After the 'agent' node runs, check the output. # tools_condition automatically checks: Did the agent output a tool_call? # - If YES: route to the "tools" node. # - If NO: route to END. workflow.add_conditional_edges( "agent", tools_condition, { "tools": "tools", # If tool call, go to tools END: "supervisor" # (CHANGED) If done with tools, go to supervisor instead of END } ) # After the tools finish executing, ALWAYS route back to the agent. # The agent needs to read the tool output and decide what to do next. workflow.add_edge("tools", "agent") workflow.add_conditional_edges("supervisor", route_supervisor, { "agent":"agent",END : END }) # --- Compile --- return workflow.compile()