Skip to content

Commit b771e9b

Browse files
committed
update memory
1 parent 90e2116 commit b771e9b

File tree

7 files changed

+950
-127
lines changed

7 files changed

+950
-127
lines changed

config/mcp_tool_definitions.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -63,16 +63,17 @@ def _get_read_code_mem_tool() -> Dict[str, Any]:
6363
"""Read code memory tool definition - reads from implement_code_summary.md"""
6464
return {
6565
"name": "read_code_mem",
66-
"description": "Check if file summary exists in implement_code_summary.md before reading actual file. Returns summary if available, otherwise recommends using read_file.",
66+
"description": "Check if file summaries exist in implement_code_summary.md for multiple files in a single call. Returns summaries for all requested files if available.",
6767
"input_schema": {
6868
"type": "object",
6969
"properties": {
70-
"file_path": {
71-
"type": "string",
72-
"description": "File path to check for summary information in implement_code_summary.md",
70+
"file_paths": {
71+
"type": "array",
72+
"items": {"type": "string"},
73+
"description": "List of file paths to check for summary information in implement_code_summary.md",
7374
}
7475
},
75-
"required": ["file_path"],
76+
"required": ["file_paths"],
7677
},
7778
}
7879

config/mcp_tool_definitions_index.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -64,16 +64,17 @@ def _get_read_code_mem_tool() -> Dict[str, Any]:
6464
"""Read code memory tool definition - reads from implement_code_summary.md"""
6565
return {
6666
"name": "read_code_mem",
67-
"description": "Check if file summary exists in implement_code_summary.md before reading actual file. Returns summary if available, otherwise recommends using read_file.",
67+
"description": "Check if file summaries exist in implement_code_summary.md for multiple files in a single call. Returns summaries for all requested files if available.",
6868
"input_schema": {
6969
"type": "object",
7070
"properties": {
71-
"file_path": {
72-
"type": "string",
73-
"description": "File path to check for summary information in implement_code_summary.md",
71+
"file_paths": {
72+
"type": "array",
73+
"items": {"type": "string"},
74+
"description": "List of file paths to check for summary information in implement_code_summary.md",
7475
}
7576
},
76-
"required": ["file_path"],
77+
"required": ["file_paths"],
7778
},
7879
}
7980

tools/code_implementation_server.py

Lines changed: 71 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import io
2020
from pathlib import Path
2121
import re
22-
from typing import Dict, Any
22+
from typing import Dict, Any, List
2323
import tempfile
2424
import shutil
2525
import logging
@@ -418,22 +418,25 @@ async def execute_bash(command: str, timeout: int = 30) -> str:
418418

419419

420420
@mcp.tool()
421-
async def read_code_mem(file_path: str) -> str:
421+
async def read_code_mem(file_paths: List[str]) -> str:
422422
"""
423-
Check if file summary exists in implement_code_summary.md
423+
Check if file summaries exist in implement_code_summary.md for multiple files
424424
425425
Args:
426-
file_path: File path to check for summary information in implement_code_summary.md
426+
file_paths: List of file paths to check for summary information in implement_code_summary.md
427427
428428
Returns:
429-
Summary information if available
429+
Summary information for all requested files if available
430430
"""
431431
try:
432-
if not file_path:
433-
result = {"status": "error", "message": "file_path parameter is required"}
434-
log_operation("read_code_mem_error", {"error": "missing_file_path"})
432+
if not file_paths or not isinstance(file_paths, list):
433+
result = {"status": "error", "message": "file_paths parameter is required and must be a list"}
434+
log_operation("read_code_mem_error", {"error": "missing_or_invalid_file_paths"})
435435
return json.dumps(result, ensure_ascii=False, indent=2)
436436

437+
# Remove duplicates while preserving order
438+
unique_file_paths = list(dict.fromkeys(file_paths))
439+
437440
# Ensure workspace exists
438441
ensure_workspace_exists()
439442

@@ -444,12 +447,12 @@ async def read_code_mem(file_path: str) -> str:
444447
if not summary_file_path.exists():
445448
result = {
446449
"status": "no_summary",
447-
"file_path": file_path,
450+
"file_paths": unique_file_paths,
448451
"message": "No summary file found.",
449-
# "recommendation": f"read_file(file_path='{file_path}')"
452+
"results": []
450453
}
451454
log_operation(
452-
"read_code_mem", {"file_path": file_path, "status": "no_summary_file"}
455+
"read_code_mem", {"file_paths": unique_file_paths, "status": "no_summary_file"}
453456
)
454457
return json.dumps(result, ensure_ascii=False, indent=2)
455458

@@ -460,54 +463,78 @@ async def read_code_mem(file_path: str) -> str:
460463
if not summary_content.strip():
461464
result = {
462465
"status": "no_summary",
463-
"file_path": file_path,
466+
"file_paths": unique_file_paths,
464467
"message": "Summary file is empty.",
465-
# "recommendation": f"read_file(file_path='{file_path}')"
468+
"results": []
466469
}
467470
log_operation(
468-
"read_code_mem", {"file_path": file_path, "status": "empty_summary"}
471+
"read_code_mem", {"file_paths": unique_file_paths, "status": "empty_summary"}
469472
)
470473
return json.dumps(result, ensure_ascii=False, indent=2)
471474

472-
# Extract file-specific section from summary
473-
file_section = _extract_file_section_from_summary(summary_content, file_path)
474-
475-
if file_section:
476-
result = {
477-
"status": "summary_found",
478-
"file_path": file_path,
479-
"summary_content": file_section,
480-
"message": f"Summary information found for {file_path} in implement_code_summary.md",
481-
}
482-
log_operation(
483-
"read_code_mem",
484-
{
475+
# Process each file path and collect results
476+
results = []
477+
summaries_found = 0
478+
479+
for file_path in unique_file_paths:
480+
# Extract file-specific section from summary
481+
file_section = _extract_file_section_from_summary(summary_content, file_path)
482+
483+
if file_section:
484+
file_result = {
485485
"file_path": file_path,
486486
"status": "summary_found",
487-
"section_length": len(file_section),
488-
},
489-
)
490-
return json.dumps(result, ensure_ascii=False, indent=2)
487+
"summary_content": file_section,
488+
"message": f"Summary information found for {file_path}"
489+
}
490+
summaries_found += 1
491+
else:
492+
file_result = {
493+
"file_path": file_path,
494+
"status": "no_summary",
495+
"summary_content": None,
496+
"message": f"No summary found for {file_path}"
497+
}
498+
499+
results.append(file_result)
500+
501+
# Determine overall status
502+
if summaries_found == len(unique_file_paths):
503+
overall_status = "all_summaries_found"
504+
elif summaries_found > 0:
505+
overall_status = "partial_summaries_found"
491506
else:
492-
result = {
493-
"status": "no_summary",
494-
"file_path": file_path,
495-
"message": f"No summary found for {file_path} in implement_code_summary.md",
496-
# "recommendation": f"Use read_file tool to read the actual file: read_file(file_path='{file_path}')"
497-
}
498-
log_operation(
499-
"read_code_mem", {"file_path": file_path, "status": "no_match"}
500-
)
501-
return json.dumps(result, ensure_ascii=False, indent=2)
507+
overall_status = "no_summaries_found"
508+
509+
result = {
510+
"status": overall_status,
511+
"file_paths": unique_file_paths,
512+
"total_requested": len(unique_file_paths),
513+
"summaries_found": summaries_found,
514+
"message": f"Found summaries for {summaries_found}/{len(unique_file_paths)} files",
515+
"results": results
516+
}
517+
518+
log_operation(
519+
"read_code_mem",
520+
{
521+
"file_paths": unique_file_paths,
522+
"status": overall_status,
523+
"total_requested": len(unique_file_paths),
524+
"summaries_found": summaries_found,
525+
},
526+
)
527+
528+
return json.dumps(result, ensure_ascii=False, indent=2)
502529

503530
except Exception as e:
504531
result = {
505532
"status": "error",
506533
"message": f"Failed to check code memory: {str(e)}",
507-
"file_path": file_path,
508-
# "recommendation": "Use read_file tool instead"
534+
"file_paths": file_paths if isinstance(file_paths, list) else [str(file_paths)],
535+
"results": []
509536
}
510-
log_operation("read_code_mem_error", {"file_path": file_path, "error": str(e)})
537+
log_operation("read_code_mem_error", {"file_paths": file_paths, "error": str(e)})
511538
return json.dumps(result, ensure_ascii=False, indent=2)
512539

513540

workflows/agents/code_implementation_agent.py

Lines changed: 29 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -297,9 +297,9 @@ async def _handle_read_file_with_memory_optimization(self, tool_call: Dict) -> D
297297
should_use_summary = False
298298
if self.memory_agent and self.mcp_agent:
299299
try:
300-
# Use read_code_mem MCP tool to check if summary exists
300+
# Use read_code_mem MCP tool to check if summary exists (pass file path as list)
301301
read_code_mem_result = await self.mcp_agent.call_tool(
302-
"read_code_mem", {"file_path": file_path}
302+
"read_code_mem", {"file_paths": [file_path]}
303303
)
304304

305305
# Parse the result to check if summary was found
@@ -308,8 +308,10 @@ async def _handle_read_file_with_memory_optimization(self, tool_call: Dict) -> D
308308
if isinstance(read_code_mem_result, str):
309309
try:
310310
result_data = json.loads(read_code_mem_result)
311+
# Check if any summaries were found in the results
311312
should_use_summary = (
312-
result_data.get("status") == "summary_found"
313+
result_data.get("status") in ["all_summaries_found", "partial_summaries_found"]
314+
and result_data.get("summaries_found", 0) > 0
313315
)
314316
except json.JSONDecodeError:
315317
should_use_summary = False
@@ -323,7 +325,7 @@ async def _handle_read_file_with_memory_optimization(self, tool_call: Dict) -> D
323325
# Use the MCP agent to call read_code_mem tool
324326
if self.mcp_agent:
325327
result = await self.mcp_agent.call_tool(
326-
"read_code_mem", {"file_path": file_path}
328+
"read_code_mem", {"file_paths": [file_path]}
327329
)
328330

329331
# Modify the result to indicate it was originally a read_file call
@@ -334,9 +336,25 @@ async def _handle_read_file_with_memory_optimization(self, tool_call: Dict) -> D
334336
json.loads(result) if isinstance(result, str) else result
335337
)
336338
if isinstance(result_data, dict):
337-
result_data["original_tool"] = "read_file"
338-
result_data["optimization"] = "redirected_to_read_code_mem"
339-
final_result = json.dumps(result_data, ensure_ascii=False)
339+
# Extract the specific file result for the single file we requested
340+
file_results = result_data.get("results", [])
341+
if file_results and len(file_results) > 0:
342+
specific_result = file_results[0] # Get the first (and only) result
343+
# Transform to match the old single-file format for backward compatibility
344+
transformed_result = {
345+
"status": specific_result.get("status", "no_summary"),
346+
"file_path": specific_result.get("file_path", file_path),
347+
"summary_content": specific_result.get("summary_content"),
348+
"message": specific_result.get("message", ""),
349+
"original_tool": "read_file",
350+
"optimization": "redirected_to_read_code_mem"
351+
}
352+
final_result = json.dumps(transformed_result, ensure_ascii=False)
353+
else:
354+
# Fallback if no results
355+
result_data["original_tool"] = "read_file"
356+
result_data["optimization"] = "redirected_to_read_code_mem"
357+
final_result = json.dumps(result_data, ensure_ascii=False)
340358
else:
341359
final_result = result
342360
except (json.JSONDecodeError, TypeError):
@@ -914,7 +932,7 @@ async def test_summary_functionality(self, test_file_path: str = None):
914932
if self.mcp_agent:
915933
try:
916934
result = await self.mcp_agent.call_tool(
917-
"read_code_mem", {"file_path": file_path}
935+
"read_code_mem", {"file_paths": [file_path]}
918936
)
919937

920938
# Parse the result to check if summary was found
@@ -924,7 +942,7 @@ async def test_summary_functionality(self, test_file_path: str = None):
924942
json.loads(result) if isinstance(result, str) else result
925943
)
926944

927-
if result_data.get("status") == "summary_found":
945+
if result_data.get("status") in ["all_summaries_found", "partial_summaries_found"] and result_data.get("summaries_found", 0) > 0:
928946
summary_files_found += 1
929947
except Exception as e:
930948
self.logger.warning(
@@ -1016,15 +1034,15 @@ async def test_summary_optimization(self, test_file_path: str = "rice/config.py"
10161034
try:
10171035
# Use MCP agent to call read_code_mem tool
10181036
result = await self.mcp_agent.call_tool(
1019-
"read_code_mem", {"file_path": test_file_path}
1037+
"read_code_mem", {"file_paths": [test_file_path]}
10201038
)
10211039

10221040
# Parse the result to check if summary was found
10231041
import json
10241042

10251043
result_data = json.loads(result) if isinstance(result, str) else result
10261044

1027-
return result_data.get("status") == "summary_found"
1045+
return result_data.get("status") in ["all_summaries_found", "partial_summaries_found"] and result_data.get("summaries_found", 0) > 0
10281046
except Exception as e:
10291047
self.logger.warning(f"Failed to test read_code_mem optimization: {e}")
10301048
return False

0 commit comments

Comments
 (0)