Skip to content

Commit 710d95e

Browse files
avimallickavimallick
committed
Final changes for explainer agent
1 parent 4b2c453 commit 710d95e

10 files changed

Lines changed: 240 additions & 81 deletions

File tree

‎AntiPattern_Remediator/main.py‎

Lines changed: 50 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
1-
21
"""
32
Main entry point - Legacy Code Migration Tool
43
"""
54
from config.settings import initialize_settings
6-
# from scripts import seed_database
75
from dotenv import load_dotenv
86
load_dotenv()
97
from colorama import Fore, Style
@@ -12,11 +10,34 @@
1210
from full_repo_workflow import run_full_repo_workflow
1311

1412

15-
def run_code_snippet_workflow(settings, db_manager, prompt_manager, langgraph):
16-
"""Run the original workflow with a hardcoded Java code snippet."""
17-
print(Fore.BLUE + "\n=== Code Snippet Analysis Workflow ===" + Style.RESET_ALL)
18-
print("Analyzing the provided Java code snippet...")
19-
13+
def main():
14+
"""Main function: Run antipattern analysis"""
15+
16+
# Let user select provider
17+
print("Available providers: 1) ollama 2) ibm 3) vllm")
18+
choice = input("Select provider (1-3): ").strip()
19+
20+
provider_map = {"1": "ollama", "2": "ibm", "3": "vllm"}
21+
provider = provider_map.get(choice, "ollama") # default to ollama
22+
23+
# Let us choose which DB to interact with
24+
print("Choose your trove: 1) ChromaDB (VectorDB) 2) TinyDB (DocumentDB)")
25+
db_choice = input("Choose 1 or 2: ").strip()
26+
27+
# Initialize global settings with selected provider
28+
settings = initialize_settings(provider)
29+
print(Fore.GREEN + f"Using {settings.LLM_PROVIDER} with model {settings.LLM_MODEL}" + Style.RESET_ALL)
30+
31+
# Temporary Lazy Imports
32+
from src.core.graph import CreateGraph
33+
from src.data.database import VectorDBManager, TinyDBManager
34+
from src.core.prompt import PromptManager
35+
from scripts import seed_database
36+
37+
# Initialize PromptManager
38+
print("Initializing PromptManager...")
39+
prompt_manager = PromptManager()
40+
2041
# Example Java code
2142
legacy_code = """
2243
public class ApplicationManager {
@@ -55,7 +76,8 @@ def run_code_snippet_workflow(settings, db_manager, prompt_manager, langgraph):
5576
}
5677
}
5778
"""
58-
79+
80+
# Initial workflow state
5981
initial_state = {
6082
"code": legacy_code,
6183
"context": None,
@@ -66,58 +88,12 @@ def run_code_snippet_workflow(settings, db_manager, prompt_manager, langgraph):
6688
"code_review_results": None,
6789
"code_review_times": 0,
6890
"msgs": [],
69-
"answer": None
70-
}
71-
72-
final_state = langgraph.invoke(initial_state)
91+
"answer": None,
7392

74-
print(Fore.GREEN + f"\nAnalysis Complete!" + Style.RESET_ALL)
75-
print(f"Final state keys: {list(final_state.keys())}")
76-
print(f"Context retrieved: {'Yes' if final_state.get('context') else 'No'}")
77-
print(f"Analysis completed: {'Yes' if final_state.get('antipatterns_scanner_results') else 'No'}")
78-
print(f"Refactored code: {'Yes' if final_state.get('refactored_code') else 'No'}")
79-
print(f"Code review results: {final_state.get('code_review_times')}")
80-
81-
82-
def main():
83-
"""Main function: Choose between code snippet analysis or full repository run"""
84-
85-
print(Fore.BLUE + "=== AntiPattern Remediator Tool ===" + Style.RESET_ALL)
86-
print("Choose your analysis mode:")
87-
print("1) Code Snippet Analysis - Analyze a sample Java code snippet")
88-
print("2) Full Repository Run - Process files with 100% test coverage from JaCoCo results")
89-
90-
# Let user choose analysis mode
91-
mode_choice = input("\nSelect mode (1-2): ").strip()
92-
93-
if mode_choice not in ["1", "2"]:
94-
print(Fore.RED + "Invalid choice. Defaulting to Code Snippet Analysis." + Style.RESET_ALL)
95-
mode_choice = "1"
96-
97-
# Let user select provider
98-
print("\nAvailable providers: 1) ollama 2) ibm 3) vllm")
99-
choice = input("Select provider (1-3): ").strip()
100-
101-
provider_map = {"1": "ollama", "2": "ibm", "3": "vllm"}
102-
provider = provider_map.get(choice, "ollama") # default to ollama
103-
104-
# Let us choose which DB to interact with
105-
print("Choose your trove: 1) ChromaDB (VectorDB) 2) TinyDB (DocumentDB)")
106-
db_choice = input("Choose 1 or 2: ").strip()
107-
108-
# Initialize global settings with selected provider
109-
settings = initialize_settings(provider)
110-
print(Fore.GREEN + f"Using {settings.LLM_PROVIDER} with model {settings.LLM_MODEL}" + Style.RESET_ALL)
111-
112-
# Temporary Lazy Imports
113-
from src.core.graph import CreateGraph
114-
from src.data.database import VectorDBManager, TinyDBManager
115-
from src.core.prompt import PromptManager
116-
from scripts import seed_database
117-
118-
# Initialize PromptManager
119-
print("Initializing PromptManager...")
120-
prompt_manager = PromptManager()
93+
# ExplainerAgent fields
94+
"explanation_response_raw": None,
95+
"explanation_json": None,
96+
}
12197

12298
# Setup Database
12399
if db_choice == "2":
@@ -133,11 +109,22 @@ def main():
133109
retriever = db_manager.as_retriever()
134110
langgraph = CreateGraph(db_manager, prompt_manager, retriever=retriever).workflow
135111

136-
# Run the selected workflow
137-
if mode_choice == "1":
138-
run_code_snippet_workflow(settings, db_manager, prompt_manager, langgraph)
112+
# Final results summary
113+
print(Fore.GREEN + f"\nAnalysis Complete!" + Style.RESET_ALL)
114+
print(f"Final state keys: {list(final_state.keys())}")
115+
print(f"Context retrieved: {'Yes' if final_state.get('context') else 'No'}")
116+
print(f"Analysis completed: {'Yes' if final_state.get('antipatterns_scanner_results') else 'No'}")
117+
print(f"Refactored code: {'Yes' if final_state.get('refactored_code') else 'No'}")
118+
print(f"Code review results: {final_state.get('code_review_times')}")
119+
120+
# Show explanation from ExplainerAgent
121+
if final_state.get("explanation_json"):
122+
import json
123+
print(Fore.CYAN + "\n=== Explanation (JSON) ===" + Style.RESET_ALL)
124+
print(json.dumps(final_state["explanation_json"], indent=2, ensure_ascii=False))
139125
else:
140-
run_full_repo_workflow(settings, db_manager, prompt_manager, langgraph)
126+
print(Fore.RED + "\nNo explanation was generated." + Style.RESET_ALL)
127+
141128

142129
if __name__ == "__main__":
143130
main()

‎AntiPattern_Remediator/src/core/agents/__init__.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,12 @@
77
from .refactor_strategist import RefactorStrategist
88
from .code_transformer import CodeTransformer
99
from .code_reviewer import CodeReviewerAgent
10+
from .explainer import ExplainerAgent
1011

1112
__all__ = [
1213
"AntipatternScanner",
1314
"RefactorStrategist",
1415
"CodeTransformer",
15-
"CodeReviewerAgent"
16+
"CodeReviewerAgent",
17+
"ExplainerAgent"
1618
]

‎AntiPattern_Remediator/src/core/agents/explainer.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
"""
2+
ExplainerAgent — minimal version
3+
- Delegates state handling to create_graph.py
4+
- Only builds messages and parses JSON response
5+
- Keeps code minimal and focused
6+
"""
7+
from __future__ import annotations
8+
9+
from typing import Dict, Any
10+
import json
11+
12+
from langchain_core.language_models import BaseLanguageModel
13+
from ..prompt import PromptManager
14+
from src.core.utils import extract_first_json
15+
16+
PROMPT_KEY = "explainer"
17+
18+
19+
class ExplainerAgent:
20+
def __init__(self, llm: BaseLanguageModel, prompt_manager: PromptManager):
21+
self.llm = llm
22+
self.prompt_manager = prompt_manager
23+
24+
def explain_antipattern(self, state: Dict[str, Any]) -> Dict[str, Any]:
25+
"""Generate explanation JSON for detected antipatterns and refactor."""
26+
kwargs = dict(
27+
code=state.get("code", ""),
28+
language=state.get("language", "Java"),
29+
context=state.get("context", ""),
30+
refactored_code=state.get("refactored_code", ""),
31+
refactor_rationale=state.get("refactor_rationale", ""),
32+
antipattern_name=state.get("antipattern_name", "Unknown antipattern"),
33+
antipattern_description=state.get("antipattern_description", ""),
34+
antipatterns_json=json.dumps(state.get("antipatterns_json", []), ensure_ascii=False),
35+
)
36+
37+
messages = self._build_messages(**kwargs)
38+
response = self.llm.invoke(messages)
39+
raw = getattr(response, "content", None) or str(response)
40+
state["explanation_response_raw"] = raw
41+
42+
parsed = extract_first_json(raw)
43+
state["explanation_json"] = parsed if isinstance(parsed, dict) else {}
44+
return state
45+
46+
def display_explanation(self, state: Dict[str, Any]) -> Dict[str, Any]:
47+
print("\n=== Explanation (raw) ===\n", state.get("explanation_response_raw", "N/A"))
48+
if state.get("explanation_json"):
49+
print("\n=== Explanation (JSON) ===\n", json.dumps(state["explanation_json"], indent=2, ensure_ascii=False))
50+
return state
51+
52+
def _build_messages(self, **kwargs) -> Any:
53+
try:
54+
getp = getattr(self.prompt_manager, "get_prompt", None)
55+
if callable(getp):
56+
prompt = getp(PROMPT_KEY)
57+
if prompt is not None:
58+
return prompt.format_messages(**kwargs)
59+
except Exception:
60+
pass
61+
62+
schema = {
63+
"items": [{
64+
"antipattern_name": "",
65+
"antipattern_description": "",
66+
"impact": "",
67+
"why_it_is_bad": "",
68+
"how_we_fixed_it": "",
69+
"refactored_code": "",
70+
"summary": ""
71+
}],
72+
"what_changed": [],
73+
"why_better": [],
74+
"principles_applied": [],
75+
"trade_offs": [],
76+
"closing_summary": ""
77+
}
78+
content = (
79+
"Given inputs (JSON):\n" + json.dumps(kwargs, ensure_ascii=False) +
80+
"\nRespond with STRICT JSON using exactly this schema:\n" +
81+
json.dumps(schema, ensure_ascii=False)
82+
)
83+
return [
84+
{"role": "system", "content": "Return STRICT JSON only. No commentary."},
85+
{"role": "user", "content": content},
86+
]

‎AntiPattern_Remediator/src/core/graph/create_graph.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from ..agents import RefactorStrategist
1515
from ..agents import CodeTransformer
1616
from ..agents import CodeReviewerAgent
17+
from ..agents import ExplainerAgent
1718

1819
# Imports for LangSmith tracing
1920
import os
@@ -62,7 +63,7 @@ def __init__(self, db_manager, prompt_manager: PromptManager, retriever=None, ll
6263
self.prompt_manager = prompt_manager
6364
self.conditional_edges = ConditionalEdges()
6465

65-
# ✅ assign the instance attribute before use
66+
# assign the instance attribute before use
6667
self.retriever = retriever or self.db_manager.as_retriever()
6768

6869
retriever_tool = create_retriever_tool(
@@ -78,13 +79,13 @@ def __init__(self, db_manager, prompt_manager: PromptManager, retriever=None, ll
7879
"strategist": RefactorStrategist(self.llm, self.prompt_manager, retriever=self.retriever),
7980
"transformer": CodeTransformer(self.llm, self.prompt_manager),
8081
"reviewer": CodeReviewerAgent(self.llm, self.prompt_manager),
82+
"explainer": ExplainerAgent(self.llm, self.prompt_manager)
8183
}
8284

8385
# Build the LangGraph workflow
8486
self.workflow = self._build_graph()
8587

8688
def _build_graph(self):
87-
"""Build LangGraph workflow"""
8889
graph = StateGraph(AgentState)
8990

9091
# Scanner: retrieve + analyze
@@ -104,6 +105,10 @@ def _build_graph(self):
104105
graph.add_node("review_code", self.agents["reviewer"].review_code)
105106
graph.add_node("display_code_review_results", self.agents["reviewer"].display_code_review_results)
106107

108+
# Explainer: final storytelling
109+
graph.add_node("explain_antipattern", self.agents["explainer"].explain_antipattern)
110+
graph.add_node("display_explanation", self.agents["explainer"].display_explanation)
111+
107112
# Topology
108113
graph.set_entry_point("retrieve_context")
109114
graph.add_edge("retrieve_context", "analyze_antipatterns")
@@ -123,6 +128,8 @@ def _build_graph(self):
123128
},
124129
)
125130

126-
graph.add_edge("display_code_review_results", END)
131+
graph.add_edge("display_code_review_results", "explain_antipattern")
132+
graph.add_edge("explain_antipattern", "display_explanation")
133+
graph.add_edge("display_explanation", END)
127134

128135
return graph.compile()

‎AntiPattern_Remediator/src/core/prompt/prompt_manager.py‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -7,27 +7,26 @@
77
class PromptManager:
88
"""Manager for handling prompt templates and configurations."""
99
def __init__(self):
10-
# Prompt key constants, **same as YAML filenames**
10+
# Prompt key constants, **same as YAML filenames (without .yaml)**
1111
self.ANTIPATTERN_SCANNER = "antipattern_scanner"
12-
self.REFACTOR_STRATEGIST = "refactor_strategist"
12+
self.REFACTOR_STRATEGIST = "refactor_strategist"
1313
self.CODE_TRANSFORMER = "code_transformer"
1414
self.CODE_REVIEWER = "code_reviewer"
15+
self.EXPLAINER_AGENT = "explainer"
1516

1617
self.prompt_directory = settings.PROMPT_DIR
17-
# Initialize storage for prompt templates
1818
self._prompt_cache = {}
19-
# Load prompts on initialization
2019
self._load_all_prompts()
2120

2221
def _load_all_prompts(self) -> None:
2322
"""Load all prompt configurations from YAML files."""
2423
try:
25-
# Get all prompt constants and load corresponding files
2624
prompt_constants = [
2725
self.ANTIPATTERN_SCANNER,
2826
self.REFACTOR_STRATEGIST,
2927
self.CODE_TRANSFORMER,
3028
self.CODE_REVIEWER,
29+
self.EXPLAINER_AGENT,
3130
]
3231

3332
for prompt_key in prompt_constants:
@@ -55,12 +54,17 @@ def _load_prompt_from_yaml(self, filename: str, prompt_key: str) -> None:
5554
return
5655

5756
prompt_config = config[prompt_key]
58-
# Create ChatPromptTemplate
59-
self._prompt_cache[prompt_key] = ChatPromptTemplate([
60-
("system", prompt_config.get('system', '')),
61-
("user", prompt_config.get('user', '')),
62-
MessagesPlaceholder("msgs")
63-
])
57+
58+
# Build messages in (role, content) format
59+
messages = []
60+
if prompt_config.get("system"):
61+
messages.append(("system", prompt_config["system"]))
62+
if prompt_config.get("user"):
63+
messages.append(("user", prompt_config["user"]))
64+
messages.append(MessagesPlaceholder("msgs"))
65+
66+
# Use the correct constructor
67+
self._prompt_cache[prompt_key] = ChatPromptTemplate.from_messages(messages)
6468
print(f"Loaded prompt '{prompt_key}' from {filename}")
6569

6670
except Exception as e:
@@ -70,5 +74,4 @@ def get_prompt(self, prompt_key: str) -> Optional[ChatPromptTemplate]:
7074
if prompt_key not in self._prompt_cache:
7175
print(f"Warning: Prompt '{prompt_key}' not found in cache")
7276
return None
73-
74-
return self._prompt_cache[prompt_key]
77+
return self._prompt_cache[prompt_key]

‎AntiPattern_Remediator/src/core/state.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,4 +17,6 @@ class AgentState(TypedDict):
1717
code_review_times: int # Number of times code has been reviewed
1818
msgs: List[Dict[str, Any]] # Message history for conversation context
1919
answer: Optional[str] # Analysis result
20+
explanation_response_raw: Optional[str] # Raw LLM output from explainer
21+
explanation_json: Optional[Dict[str, Any]] # Parsed JSON explanation
2022
current_file_path: Optional[str] # Path to the current file being processed

‎AntiPattern_Remediator/src/core/utils/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from .json_utils import extract_first_json
2+
3+
__all__ = ["extract_first_json"]

0 commit comments

Comments
 (0)