Redesign, Audit and Revamp Code & Snippet Agents + Judge Logic

[ishaanxgupta]

The current implementation of the CodeAgent, SnippetAgent, and their associated logic in the JudgeAgent needs a comprehensive audit and revamp. Several architectural gaps and technical debts have been identified that hinder reliability and scalability.

Lets have discussion on this issue first, whoever has the best solution will be assigned the issue, we can have multiple contributors on this as well.

Bounty: $20 API Credits

[LewiX-369]

Happy to tackle this issue.

/opire try

[Ramanand-Shirbhate]

Hey @ishaanxgupta, I'd love to help tackle this revamp.

I see you want to discuss solutions first before assigning. Based on the mention of architectural gaps and technical debt in the CodeAgent, SnippetAgent, and JudgeAgent, my initial audit strategy would focus on these three pillars:

Reliability (Judge Logic): Auditing the JudgeAgent's parsing and validation logic to ensure it gracefully handles unexpected LLM outputs, hallucinations, or malformed snippets without breaking the pipeline.

Scalability (Context & State): Reviewing how the CodeAgent and SnippetAgent manage state and context windows. Tightly coupled state often creates bottlenecks, so I'd look for ways to ensure they operate more modularly.

Technical Debt: Identifying tightly coupled dependencies between the agents and decoupling them into isolated, easily testable structures.

Could you point me to the specific directory or files where these agents currently live? I'll do a quick review of the current implementation and draft a concrete technical proposal for us to discuss here.

[ved015]

Happy to tackle this issue.

/opire try

Thank you for taking interest pls share your approach here

[LewiX-369]

Hi @ved015 and @ishaanxgupta,

Thanks for getting back to me! Here is my detailed technical proposal for the audit and revamp:

1. Reliable Structured Output (Addressing Parsing Brittleness)

Currently, CodeAgent and SnippetAgent rely on fragile string parsing (splitting by ```json) which leads to pipeline breaks when the LLM hallucinates markdown or extra text.
Approach: I will migrate both agents to use LangChain's .with_structured_output(Schema, include_raw=True). This enforces strict Pydantic schema validation at the model level. I'll implement a _call_model_structured wrapper in BaseAgent to handle this cleanly while preserving access to the raw AIMessage so that our metrics tracking (_track_llm_call) remains fully intact.

2. Deterministic Judging & Deduplication (Addressing Scalability)

Currently, JudgeAgent passes all incoming items to the LLM to decide on ADD, UPDATE, DELETE, or NOOP. For highly structured domains like code and snippet, this is an expensive and unnecessary bottleneck.
Approach:

• I will introduce _deterministic_code and _deterministic_snippet methods in JudgeAgent.
• Code: Deduplicate deterministically by matching target_symbol/target_file + annotation_type. If they match and content is identical, it's a NOOP. If content differs, it's an UPDATE.
• Snippet: Exact string match on the natural language description or code for NOOP/UPDATE.
• This bypasses the LLM entirely for these domains, saving tokens, latency, and eliminating hallucinations during the judging phase.

3. Better Context for Fallback Judging (Addressing Context State)

For cases where the semantic LLM judge must be used, passing raw dictionaries is suboptimal.
Approach: I will add _code_item_to_string and _snippet_item_to_string to format the context cleanly (e.g., [bug_report] auth_service.py: Token expires too early). I'll also update src/prompts/judge.py to explicitly teach the judge the domain-specific rules for code and snippet.

I have already audited the codebase (src/agents/code.py, src/agents/snippet.py, src/agents/judge.py, src/prompts/judge.py) and validated this approach locally. I can open a PR shortly if you approve this direction!

[Ankit-Kotnala]

@ved015 @ishaanxgupta
I’d like to take this up.

Plan:

• Audit CodeAgent, SnippetAgent, and JudgeAgent flow.
• Add structured Pydantic output support in src/agents/base.py, used by:
  • src/agents/code.py
  • src/agents/snippet.py
• Add deterministic judge logic for code/snippet in src/agents/judge.py:
  • code: match by target symbol/file + annotation type
  • snippet: match by exact code or exact description
• Update src/pipelines/ingest.py to send structured code/snippet items to Judge.
• Update src/pipelines/weaver.py to store/read structured JSON while preserving old pipe-delimited records.
• Add focused regression tests in tests/test_deterministic_memory_layer.py.

Expected result:

• Less brittle parsing.
• Fewer unnecessary LLM calls.
• Safer ADD/UPDATE/NOOP decisions.
• Existing profile/temporal logic remains unchanged.

[ved015]

hi @Ramanand-Shirbhate @vovchuklena-source @LewiX-369 @Ankit-Kotnala Thank you everyone for sharing their approaches we were thinking if you can get some insight out of other open source tools like gemini-cli, opencode and see how they handle memory management (specific to code) and draft your solution then because in this issue we also want to change the complete workflow of our code agent not just refactoring , error handling and handling some edge cases we are thinking of a complete new approach (cannot just blindly store all code ofc)

Would love to hear your new approaches pls try to also share a diagram or some visual display of your arch too if possible
Xmem will be the first memory layer that handles code related stuff too 😁

[Ramanand-Shirbhate]

Hey @ved015, that is a fantastic pivot. You hit the nail on the head—blindly storing raw code stringified in a vector DB leads to massive context bloat and terrible retrieval accuracy.

Looking at tools like gemini-cli and advanced code-RAG systems, the gold standard for code memory is Hierarchical Semantic Chunking. Instead of treating code as plain text, the CodeAgent should act as a lightweight static analyzer before anything hits the memory layer.

Here is my proposed workflow redesign for XMem's code handling:

The New Workflow
AST Extraction (CodeAgent): Instead of storing the whole file, the agent uses Tree-sitter (or Python's ast) to break the code into semantic nodes (Classes, Functions, Dependencies).

Skeleton vs. Implementation: We split the data. The 'Skeleton' (Function signatures + Docstrings) is stored for fast, cheap retrieval. The 'Implementation' (the actual code body) is stored separately and linked via metadata.

Hierarchical Judging (JudgeAgent): The Judge first retrieves the Skeleton. If the LLM determines it needs the exact logic to answer the prompt, it fetches the linked Implementation.

Proposed Architecture Diagram

Why this solves the problem:

It drastically reduces token usage because the Judge operates mostly on Skeletons.

It prevents "blind storage" by giving the system an actual structural understanding of the code.

If this aligns with the vision for XMem's new memory layer, I'd love to take point on implementing the CodeAgent AST parser to get this rolling!

[ishaanxgupta]

Hi @Ramanand-Shirbhate Thanks for the approach. are you planning to give code agent tools like AST, docstring read etc.
I will give you an example the user could simple enter a code snippet/whole file and ask for its explanation or the agent could give a snippet/file as an output. how do you plan to store the snippet? Could we give our code agent 3-4 tools and according to the input code agent will handle those? what do you think about this type of an approach?
There could be a lot of edge cases in the input parameter when the code agent gets invoked and we don't have a tool to handle the input!
Could code embedding also help us in any of the use case?

[Ramanand-Shirbhate]

Hey @ishaanxgupta, spot on. A tool-calling architecture is exactly what I was leaning towards—it's definitely the cleanest way to handle unpredictable inputs without breaking the pipeline.
Here’s how I’d see those specific ideas playing out:

1. Handling diverse inputs (The Router pattern)
   Having the CodeAgent act as a smart router makes perfect sense. We give the LLM a few specific tools and let it decide what to use based on what the user throws at it:
   parse_file_ast: For when someone drops in a full .py or .js file. This does the heavy lifting of extracting skeletons and chunks.
   store_raw_snippet: For a quick 10-line paste. It skips the AST overhead, generates a quick summary, and stores it as-is.
   explain_code: For read-only queries where they just want an explanation without persisting anything to memory.

2. Storing the snippets
   For those smaller snippets, keeping it lightweight is key. The store_raw_snippet tool would just tag the language, write a 1-sentence summary of the logic, and store it as a single chunk. No need to overcomplicate it.

3. Code Embeddings
   100% yes, this is the secret sauce. Embedding massive walls of raw code usually ruins vector search accuracy. Instead, we should only embed the summaries and docstrings.

Storage: We embed the docstring/summary and drop it in the Vector DB with a metadata pointer to the actual code block.

Retrieval: If a user asks "how does the auth work?", we embed their query, run a similarity search against the docstring embeddings, and instantly pull the exact implementation it points to.

This tool-router setup combined with targeted embeddings should completely eliminate the edge-case brittleness. Does this align with where you want to take XMem?