Skip to content

Commit 9199a45

Browse files
committed
added agent fixes
Signed-off-by: Arya Pratap Singh <notaryasingh@gmail.com>
1 parent 5981db0 commit 9199a45

File tree

13 files changed

+42
-57
lines changed

13 files changed

+42
-57
lines changed
640 KB
Binary file not shown.
271 KB
Binary file not shown.
76.9 KB
Binary file not shown.
-111 Bytes
Binary file not shown.

agent/dsa_agent/agent.py

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
from dsa_agent.state import AgentState
33
from dsa_agent.nodes import (
44
retrieve_question,
5-
update_question,
65
code_generation_in_node,
76
visualize_code,
87
complexity_analysis,
@@ -21,7 +20,7 @@
2120

2221
workflow = StateGraph(AgentState)
2322

24-
workflow.add_node("update_question", update_question)
23+
# workflow.add_node("update_question", update_question)
2524
workflow.add_node("retrieve_question", retrieve_question)
2625
workflow.add_node("code_generation_in_node", code_generation_in_node)
2726
workflow.add_node("visualize_code", visualize_code)
@@ -32,67 +31,61 @@
3231
workflow.add_conditional_edges(
3332
START,
3433
new_question,
35-
{
36-
"update_question": "update_question",
37-
"retrieve_question": "retrieve_question",
38-
},
39-
)
34+
{"retrieve_question": "retrieve_question", "no_context": "no_context"},
4035

36+
)
4137

4238

4339
workflow.add_edge("retrieve_question", "code_generation_in_node")
4440

4541

46-
4742
workflow.add_conditional_edges(
4843
"code_generation_in_node",
4944
decide_to_generate_code,
5045
{
5146
"useful": "visualize_code",
52-
"code_generation_in_node": "code_generation_in_node",
47+
"no_context": "no_context",
5348
},
5449
)
5550

5651

5752

58-
workflow.add_edge("visualize_code", "retrieve_question")
53+
# workflow.add_edge("visualize_code", "retrieve_question")
5954

6055

6156
workflow.add_conditional_edges(
6257
"visualize_code",
6358
decide_to_generate_visualisation,
6459
{
6560
"useful": "complexity_analysis",
66-
"visualize_code": "visualize_code",
61+
"no_context": "no_context",
6762
},
6863
)
6964

70-
workflow.add_edge("complexity_analysis", "retrieve_question")
65+
# workflow.add_edge("complexity_analysis", "retrieve_question")
7166

7267

7368
workflow.add_conditional_edges(
7469
"complexity_analysis",
7570
decide_to_generate_complexity,
7671
{
7772
"useful": "explaination_code",
78-
"complexity_analysis": "complexity_analysis",
73+
"no_context": "no_context",
7974
},
8075
)
8176

82-
workflow.add_edge("explaination_code", "retrieve_question")
77+
# workflow.add_edge("explaination_code", "retrieve_question")
8378

8479

8580
workflow.add_conditional_edges(
8681
"explaination_code",
8782
decide_to_generate_explanation,
8883
{
8984
"useful": END,
90-
"explaination_code": "explaination_code",
85+
"no_context": "no_context",
9186
},
9287
)
9388

94-
workflow.add_edge("update_question", END)
95-
9689
workflow.add_edge("no_context", END)
9790

9891
graph = workflow.compile(checkpointer=MemorySaver())

agent/dsa_agent/edges/__init__.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,16 @@
66
from dsa_agent.edges.explanation_verify import explaination_verify
77

88
def new_question(state):
9-
messages = state["messages"]
10-
last_message = messages[-1]
9+
question = state["question"]
10+
# last_message = messages[-1]
1111

1212
if (
13-
isinstance(last_message, SystemMessage)
14-
and "QUESTION UPDATED" in last_message.content
13+
isinstance(question, SystemMessage)
14+
and "QUESTION UPDATED" in question.content
1515
):
1616
return "update_question"
17+
18+
print("init edges new question : ", question)
1719

1820
return "retrieve_question"
1921

@@ -32,6 +34,7 @@ def decide_to_generate_code(state):
3234
print("---ASSESS GENERATED CODE---")
3335
code = state["code"]
3436
question = state["question"]
37+
# testCases = state["testCases"]
3538

3639
print("---GRADE CODE---")
3740
score = code_verify.invoke({"question": question, "code": code})

agent/dsa_agent/edges/code_verify.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,13 @@ class VerifyCode(BaseModel):
1616
code_verification = llm.with_structured_output(VerifyCode)
1717

1818
system = """You a 4000 rated Competitive Programmer that writes optimised and fully detailed working code in python \n
19-
Given a question and testCases you are expected to write a working code in python that passes all the testCases and is generic and well optimised \n"""
19+
Given a question you are expected to write a working code in python that is generic and well optimised \n"""
20+
2021

2122
code_prompt = ChatPromptTemplate.from_messages(
2223
[
2324
("system", system),
24-
("human", "User question: \n\n {question} \n\n Code Generated: \n\n {code}"),
25+
("human", "User question: \n\n {question} \n\n Code Generated: \n\n {code} \n\n"),
2526
]
2627
)
2728

agent/dsa_agent/nodes/__init__.py

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,8 @@ def retrieve_question(state):
1616
state (dict): New key added to state, documents, that contains retrieved documents
1717
"""
1818
print("---GETTING QUESTION---")
19-
messages = state["messages"]
20-
question = messages[-1].content
21-
19+
question = state["question"]
20+
print(f"---nodes new QUESTION DETECTED: {question}---")
2221
return {
2322
**state,
2423
"question": question,
@@ -49,9 +48,12 @@ def code_generation_in_node(state):
4948

5049
print("---TRANSFORM QUERY---")
5150
question = state["question"]
52-
testCases = state["testCases"]
51+
# testCases = state["testCases"]
52+
# print(testCases)
53+
# if testCases == []:
54+
# testCases = ["X RETURNS Y"]
5355

54-
code_generated = code_writer.invoke({"question": question , "testCases": testCases})
56+
code_generated = code_writer.invoke({"question": question })
5557
return {**state, "code": code_generated}
5658

5759
def visualize_code(state):
@@ -85,7 +87,7 @@ def complexity_analysis(state):
8587
code = state["code"]
8688

8789
complexity = complexity_generated.invoke({"code": code})
88-
return {**state, "time_complexity": list(complexity)[0], "space_complexity": list(complexity)[1]}
90+
return {**state, "time_complexity": complexity, "space_complexity": complexity}
8991

9092
def explaination_code(state):
9193
'''
@@ -107,7 +109,7 @@ def explaination_code(state):
107109
def no_context(state):
108110
print("---NO CONTEXT---")
109111

110-
messages = state["messages"]
112+
messages = state["question"]
111113
messages.append(HumanMessage("I'm sorry, I can't find any relevant information."))
112114

113115
return state

agent/dsa_agent/nodes/coding.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,14 @@
88
llm = ChatGroq(model="llama-3.3-70b-versatile", temperature=0)
99

1010
system = """You a 4000 rated Competitive Programmer that writes optimised and fully detailed working code in python \n
11-
Given a question and testCases you are expected to write a working code in python that passes all the testCases and is generic and well optimised \n"""
11+
Given a question you are expected to write a working code in python that is generic and well optimised \n"""
1212

1313
re_write_prompt = ChatPromptTemplate.from_messages(
1414
[
1515
("system", system),
1616
(
1717
"human",
18-
"Here is the question: \n\n {question} \n and here are the test cases : \n\n {testCases}. Formulate an optimise answer in python.",
18+
"Here is the question: \n\n {question} \n. Formulate an optimised answer in python.",
1919
),
2020
]
2121
)

agent/dsa_agent/state.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,18 +9,18 @@ class AgentState(TypedDict):
99
1010
Attributes:
1111
question: str
12-
testCases: List[str]
12+
# testCases: List[str]
1313
code: Optional[str]
1414
explanation: Optional[str]
1515
time_complexity: Optional[str]
1616
space_complexity: Optional[str]
1717
visualization: Optional[str]
18-
messages: list of messages
18+
# messages: list of messages
1919
"""
2020

2121
question: str
22-
testCases: List[str]
23-
messages: Annotated[Sequence[BaseMessage], add_messages]
22+
# testCases: Optional[List[str]] # type: ignore
23+
# messages: Annotated[Sequence[BaseMessage], add_messages]
2424
code: Optional[str]
2525
explanation: Optional[str]
2626
time_complexity: Optional[str]

0 commit comments

Comments
 (0)