Skip to content

Commit 17fb8ec

Browse files
committed
[bugfix][server bound match]
1 parent d5e2232 commit 17fb8ec

File tree

6 files changed

+414
-63
lines changed

6 files changed

+414
-63
lines changed
Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
import os
2+
import json
3+
from loguru import logger
4+
5+
try:
6+
import os, sys
7+
src_dir = os.path.join(
8+
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
9+
)
10+
sys.path.append(src_dir)
11+
import test_config
12+
api_key = os.environ["OPENAI_API_KEY"]
13+
api_base_url= os.environ["API_BASE_URL"]
14+
model_name = os.environ["model_name"]
15+
embed_model = os.environ["embed_model"]
16+
embed_model_path = os.environ["embed_model_path"]
17+
except Exception as e:
18+
# set your config
19+
api_key = ""
20+
api_base_url= ""
21+
model_name = ""
22+
embed_model = ""
23+
embed_model_path = ""
24+
logger.error(f"{e}")
25+
26+
from muagent.base_configs.env_config import CB_ROOT_PATH
27+
from muagent.llm_models.llm_config import EmbedConfig, LLMConfig
28+
from muagent.connector.phase import BasePhase
29+
from muagent.connector.agents import BaseAgent, SelectorAgent
30+
from muagent.connector.chains import BaseChain
31+
from muagent.connector.schema import Message, Role, ChainConfig
32+
from muagent.codechat.codebase_handler.codebase_handler import CodeBaseHandler
33+
34+
from muagent.tools import CodeRetrievalSingle
35+
36+
37+
38+
# # 下面给定了一份代码片段,以及来自于它的依赖类、依赖方法相关的代码片段,你需要判断是否为这段指定代码片段生成测例。
39+
# # 理论上所有代码都需要写测例,但是受限于人的精力不可能覆盖所有代码
40+
# # 考虑以下因素进行裁剪:
41+
# # 功能性: 如果它实现的是一个具体的功能或逻辑,则通常需要编写测试用例以验证其正确性。
42+
# # 复杂性: 如果代码较为,尤其是包含多个条件判断、循环、异常处理等的代码,更可能隐藏bug,因此应该编写测试用例。如果代码涉及复杂的算法或者逻辑,那么编写测试用例可以帮助确保逻辑的正确性,并在未来的重构中防止引入错误。
43+
# # 关键性: 如果它是关键路径的一部分或影响到核心功能,那么它就需要被测试。对于核心业务逻辑或者系统的关键组件,应当编写全面的测试用例来确保功能的正确性和稳定性。
44+
# # 依赖性: 如果代码有外部依赖,可能需要编写集成测试或模拟这些依赖进行单元测试。
45+
# # 用户输入: 如果代码处理用户输入,尤其是来自外部的、非受控的输入,那么创建测试用例来检查输入验证和处理是很重要的。
46+
# # 频繁更改:对于经常需要更新或修改的代码,有相应的测试用例可以确保更改不会破坏现有功能。
47+
48+
49+
# # 代码公开或重用:如果代码将被公开或用于其他项目,编写测试用例可以提高代码的可信度和易用性。
50+
51+
52+
# update new agent configs
53+
judgeGenerateTests_PROMPT = """#### Agent Profile
54+
When determining the necessity of writing test cases for a given code snippet,
55+
it's essential to evaluate its interactions with dependent classes and methods (retrieved code snippets),
56+
in addition to considering these critical factors:
57+
1. Functionality: If it implements a concrete function or logic, test cases are typically necessary to verify its correctness.
58+
2. Complexity: If the code is complex, especially if it contains multiple conditional statements, loops, exceptions handling, etc.,
59+
it's more likely to harbor bugs, and thus test cases should be written.
60+
If the code involves complex algorithms or logic, then writing test cases can help ensure the accuracy of the logic and prevent errors during future refactoring.
61+
3. Criticality: If it's part of the critical path or affects core functionalities, then it needs to be tested.
62+
Comprehensive test cases should be written for core business logic or key components of the system to ensure the correctness and stability of the functionality.
63+
4. Dependencies: If the code has external dependencies, integration testing may be necessary, or mocking these dependencies during unit testing might be required.
64+
5. User Input: If the code handles user input, especially from unregulated external sources, creating test cases to check input validation and handling is important.
65+
6. Frequent Changes: For code that requires regular updates or modifications, having the appropriate test cases ensures that changes do not break existing functionalities.
66+
67+
#### Input Format
68+
69+
**Code Snippet:** the initial Code or objective that the user wanted to achieve
70+
71+
**Retrieval Code Snippets:** These are the associated code segments that the main Code Snippet depends on.
72+
Examine these snippets to understand how they interact with the main snippet and to determine how they might affect the overall functionality.
73+
74+
#### Response Output Format
75+
**Action Status:** Set to 'finished' or 'continued'.
76+
If set to 'finished', the code snippet does not warrant the generation of a test case.
77+
If set to 'continued', the code snippet necessitates the creation of a test case.
78+
79+
**REASON:** Justify the selection of 'finished' or 'continued', contemplating the decision through a step-by-step rationale.
80+
"""
81+
82+
generateTests_PROMPT = """#### Agent Profile
83+
As an agent specializing in software quality assurance,
84+
your mission is to craft comprehensive test cases that bolster the functionality, reliability, and robustness of a specified Code Snippet.
85+
This task is to be carried out with a keen understanding of the snippet's interactions with its dependent classes and methods—collectively referred to as Retrieval Code Snippets.
86+
Analyze the details given below to grasp the code's intended purpose, its inherent complexity, and the context within which it operates.
87+
Your constructed test cases must thoroughly examine the various factors influencing the code's quality and performance.
88+
89+
ATTENTION: response carefully referenced "Response Output Format" in format.
90+
91+
Each test case should include:
92+
1. clear description of the test purpose.
93+
2. The input values or conditions for the test.
94+
3. The expected outcome or assertion for the test.
95+
4. Appropriate tags (e.g., 'functional', 'integration', 'regression') that classify the type of test case.
96+
5. these test code should have package and import
97+
98+
#### Input Format
99+
100+
**Code Snippet:** the initial Code or objective that the user wanted to achieve
101+
102+
**Retrieval Code Snippets:** These are the interrelated pieces of code sourced from the codebase, which support or influence the primary Code Snippet.
103+
104+
#### Response Output Format
105+
**SaveFileName:** construct a local file name based on Question and Context, such as
106+
107+
```java
108+
package/class.java
109+
```
110+
111+
**Test Code:** generate the test code for the current Code Snippet.
112+
```java
113+
...
114+
```
115+
116+
"""
117+
118+
from muagent.tools import CodeRetrievalSingle, RelatedVerticesRetrival, Vertex2Code
119+
120+
# 定义一个新的agent类
121+
class CodeRetrieval(BaseAgent):
122+
123+
def start_action_step(self, message: Message) -> Message:
124+
'''do action before agent predict '''
125+
# 根据问题获取代码片段和节点信息
126+
action_json = CodeRetrievalSingle.run(message.code_engine_name, message.input_query, llm_config=self.llm_config, embed_config=self.embed_config, search_type="tag",
127+
local_graph_path=message.local_graph_path, use_nh=message.use_nh)
128+
current_vertex = action_json['vertex']
129+
message.customed_kargs["Code Snippet"] = action_json["code"]
130+
message.customed_kargs['Current_Vertex'] = current_vertex
131+
132+
# 获取邻近节点
133+
action_json = RelatedVerticesRetrival.run(message.code_engine_name, message.customed_kargs['Current_Vertex'])
134+
# 获取邻近节点所有代码
135+
relative_vertex = []
136+
retrieval_Codes = []
137+
for vertex in action_json["vertices"]:
138+
# 由于代码是文件级别,所以相同文件代码不再获取
139+
# logger.debug(f"{current_vertex}, {vertex}")
140+
current_vertex_name = current_vertex.replace("#", "").replace(".java", "" ) if current_vertex.endswith(".java") else current_vertex
141+
if current_vertex_name.split("#")[0] == vertex.split("#")[0]: continue
142+
143+
action_json = Vertex2Code.run(message.code_engine_name, vertex)
144+
if action_json["code"]:
145+
retrieval_Codes.append(action_json["code"])
146+
relative_vertex.append(vertex)
147+
#
148+
message.customed_kargs["Retrieval_Codes"] = retrieval_Codes
149+
message.customed_kargs["Relative_vertex"] = relative_vertex
150+
151+
code_snippet = message.customed_kargs.get("Code Snippet", "")
152+
current_vertex = message.customed_kargs.get("Current_Vertex", "")
153+
message.customed_kargs["Code Snippet"] = f"name: {current_vertex}\n{code_snippet}"
154+
155+
Retrieval_Codes = message.customed_kargs["Retrieval_Codes"]
156+
Relative_vertex = message.customed_kargs["Relative_vertex"]
157+
message.customed_kargs["Retrieval Code Snippets"] = "\n".join([
158+
f"name: {vertext}\n{code}" for vertext, code in zip(Relative_vertex, Retrieval_Codes)
159+
])
160+
161+
return message
162+
163+
164+
llm_config = LLMConfig(
165+
model_name="gpt-4", api_key=api_key, api_base_url=api_base_url, temperature=0.3
166+
)
167+
embed_config = EmbedConfig(
168+
embed_engine="model", embed_model=embed_model, embed_model_path=embed_model_path
169+
)
170+
171+
172+
# initialize codebase
173+
# delete codebase
174+
codebase_name = 'client_local'
175+
code_path = "D://chromeDownloads/devopschat-bot/client_v2/client"
176+
use_nh = True
177+
do_interpret = False
178+
# cbh = CodeBaseHandler(codebase_name, code_path, crawl_type='dir', use_nh=use_nh, local_graph_path=CB_ROOT_PATH,
179+
# llm_config=llm_config, embed_config=embed_config)
180+
# cbh.delete_codebase(codebase_name=codebase_name)
181+
182+
# # load codebase
183+
# cbh = CodeBaseHandler(codebase_name, code_path, crawl_type='dir', use_nh=use_nh, local_graph_path=CB_ROOT_PATH,
184+
# llm_config=llm_config, embed_config=embed_config)
185+
# cbh.import_code(do_interpret=do_interpret)
186+
187+
188+
# log-level,print prompt和llm predict
189+
os.environ["log_verbose"] = "1"
190+
191+
CodeJudger_role = Role(role_type="assistant", role_name="CodeJudger_role", prompt=judgeGenerateTests_PROMPT)
192+
CodeJudger = CodeRetrieval(
193+
role=CodeJudger_role,
194+
chat_turn=1,
195+
llm_config=llm_config, embed_config=embed_config,
196+
)
197+
198+
199+
generateTests_role = Role(role_type="assistant", role_name="generateTests_role", prompt=generateTests_PROMPT)
200+
generateTests = CodeRetrieval(
201+
role=generateTests_role,
202+
chat_turn=1,
203+
llm_config=llm_config, embed_config=embed_config,
204+
)
205+
206+
chain_config = ChainConfig(
207+
chain_name="code2test_chain",
208+
agents=[CodeJudger_role.role_name, generateTests_role.role_name,],
209+
chat_turn=1)
210+
211+
chain = BaseChain(
212+
chainConfig=chain_config, agents=[CodeJudger, generateTests],
213+
llm_config=llm_config, embed_config=embed_config,
214+
)
215+
216+
phase = BasePhase(
217+
phase_name="code2test_phase", chains=[chain],
218+
embed_config=embed_config, llm_config=llm_config
219+
)
220+
221+
# round-1
222+
# 根据前面的load过程进行初始化
223+
cbh = CodeBaseHandler(codebase_name, code_path, crawl_type='dir', use_nh=use_nh, local_graph_path=CB_ROOT_PATH,
224+
llm_config=llm_config, embed_config=embed_config)
225+
vertexes = cbh.search_vertices(vertex_type="class")
226+
logger.debug(vertexes)
227+
228+
test_cases = []
229+
for vertex in vertexes:
230+
query_content = f"为{vertex}生成可执行的测例 "
231+
query = Message(
232+
role_name="human", role_type="user", input_query=query_content,
233+
code_engine_name=codebase_name, score_threshold=1.0, top_k=3, cb_search_type="tag",
234+
use_nh=use_nh, local_graph_path=CB_ROOT_PATH,
235+
)
236+
output_message, output_memory = phase.step(query, reinit_memory=True)
237+
# print(output_memory.to_str_messages(return_all=True, content_key="parsed_output_list"))
238+
print(output_memory.get_spec_parserd_output())
239+
values = output_memory.get_spec_parserd_output()
240+
test_code = {k:v for i in values for k,v in i.items() if k in ["SaveFileName", "Test Code"]}
241+
test_cases.append(test_code)
242+
243+
os.makedirs(f"{CB_ROOT_PATH}/tests", exist_ok=True)
244+
if "SaveFileName" in test_code:
245+
with open(f"{CB_ROOT_PATH}/tests/{test_code['SaveFileName']}", "w") as f:
246+
f.write(test_code["Test Code"])
247+
break

muagent/connector/memory_manager.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -521,6 +521,7 @@ def get_vbname_from_chatindex(self, chat_index: str) -> str:
521521
TextField("role_type", ),
522522
TextField('input_query'),
523523
TextField("role_content", ),
524+
TextField("role_tags"),
524525
TextField("parsed_output"),
525526
TextField("customed_kargs",),
526527
TextField("db_docs",),
@@ -561,7 +562,7 @@ def __init__(
561562
self.save_message_keys = [
562563
'chat_index', 'message_index', 'user_name', 'role_name', 'role_type', 'input_query', 'role_content', 'step_content',
563564
'parsed_output', 'parsed_output_list', 'customed_kargs', "db_docs", "code_docs", "search_docs", 'start_datetime', 'end_datetime',
564-
"keyword", "vector",
565+
"keyword", "vector", "role_tags"
565566
]
566567
self.use_vector = use_vector
567568
self.init_tb()
@@ -624,6 +625,9 @@ def append_tools(self, tool_information: dict, chat_index: str, nodeid: str, use
624625
pass
625626
self.append(message)
626627

628+
def get_memory_by_tag(self, tag: str) -> Memory:
629+
return self.get_memory_pool_by_key_content(key='tag', content=f'*{tag}*')
630+
627631
def get_memory_pool(self, chat_index: str = "") -> Memory:
628632
return self.get_memory_pool_by_all({"chat_index": chat_index})
629633

muagent/connector/schema/message.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ class Message(BaseModel):
1515
input_query: str = ""
1616
start_datetime: str = None
1717
end_datetime: str = None
18+
role_tags: str = ""
1819

1920
# llm output
2021
role_content: str = ""

muagent/db_handler/graph_db_handler/geabase_handler.py

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,6 @@ def get_current_nodes(self, attributes: dict, node_type: str = None, return_keys
170170

171171
nodes = result.get("n0", []) or result.get("n0.attr", [])
172172
return self.convert2GNodes(nodes)
173-
return [GNode(id=node["id"], type=node["type"], attributes=node) for node in nodes]
174173

175174
def get_current_edge(self, src_id, dst_id, edge_type:str = None, return_keys: list = []) -> GEdge:
176175
# todo 业务逻辑
@@ -183,7 +182,6 @@ def get_current_edge(self, src_id, dst_id, edge_type:str = None, return_keys: li
183182

184183
edges = result.get("e", []) or result.get("e.attr", [])
185184
return self.convert2GEdges(edges)[0]
186-
return [GEdge(start_id=edge["start_id"], end_id=edge["end_id"], type=edge["type"], attributes=edge) for edge in edges][0]
187185

188186
def get_neighbor_nodes(self, attributes: dict, node_type: str = None, return_keys: list = [], reverse=False) -> List[GNode]:
189187
#
@@ -213,7 +211,6 @@ def get_neighbor_edges(self, attributes: dict, node_type: str = None, return_key
213211

214212
edges = result.get("e", []) or result.get("e.attr", [])
215213
return self.convert2GEdges(edges)
216-
return [GEdge(start_id=edge["start_id"], end_id=edge["end_id"], type=edge["type"], attributes=edge) for edge in edges]
217214

218215
def check_neighbor_exist(self, attributes: dict, node_type: str = None, check_attributes: dict = {}) -> bool:
219216
result = self.get_neighbor_nodes(attributes, node_type,)
@@ -255,10 +252,8 @@ def get_hop_infos(self, attributes: dict, node_type: str = None, hop: int = 2, b
255252
hop -= hop_max
256253
iter_index += 1
257254

258-
nodes = self.convert2GNodes(result.get("n1", []))
255+
nodes = self.convert2GNodes(result.get("n1", [])+result.get("n0", []))
259256
edges = self.convert2GEdges(result.get("e", []))
260-
# nodes = [GNode(id=node["id"], type=node["type"], attributes=node) for node in result.get("n1", [])]
261-
# edges = [GEdge(start_id=edge["start_id"], end_id=edge["end_id"], type=edge["type"], attributes=edge) for edge in result.get("e", [])]
262257
return Graph(nodes=nodes, edges=edges, paths=result.get("p", []))
263258

264259
def get_hop_nodes(self, attributes: dict, node_type: str = None, hop: int = 2, block_attributes: dict = []) -> List[GNode]:
@@ -385,10 +380,6 @@ def decode_path(self, col_data, k) -> List:
385380
connections = {}
386381
for step in steps:
387382
props = step["props"]
388-
# if path == []:
389-
# path.append(props["original_src_id1__"].get("strVal", "") or props["original_src_id1__"].get("intVal", -1))
390-
# path.append(props["original_dst_id2__"].get("strVal", "") or props["original_dst_id2__"].get("intVal", -1))
391-
392383
start = props["original_src_id1__"].get("strVal", "") or props["original_src_id1__"].get("intVal", -1)
393384
end = props["original_dst_id2__"].get("strVal", "") or props["original_dst_id2__"].get("intVal", -1)
394385
connections[start] = end
@@ -449,10 +440,13 @@ def get_nodetypes_by_edgetype(self, edge_type: str):
449440

450441
def convert2GNodes(self, raw_nodes: List[Dict]) -> List[GNode]:
451442
nodes = []
443+
nodeids_set = set()
452444
for node in raw_nodes:
453445
node_id = node.pop("id")
454446
node_type = node.pop("type")
455-
nodes.append(GNode(id=node_id, type=node_type, attributes=node))
447+
if node_id not in nodeids_set:
448+
nodes.append(GNode(id=node_id, type=node_type, attributes=node))
449+
nodeids_set.add(node_id)
456450
return nodes
457451

458452
def convert2GEdges(self, raw_edges: List[Dict]) -> List[GEdge]:

muagent/schemas/ekg/ekg_graph.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,8 @@ class NodeSchema(BaseModel):
3131

3232
def attributes(self, ):
3333
attrs = copy.deepcopy(vars(self))
34-
for k in ["ID", "type", "id"]:
34+
# for k in ["ID", "type", "id"]:
35+
for k in ["type", "id"]:
3536
attrs.pop(k)
3637
attrs.update(json.loads(attrs.pop("extra", '{}') or '{}'))
3738
return attrs
@@ -51,7 +52,8 @@ class EdgeSchema(BaseModel):
5152

5253
def attributes(self, ):
5354
attrs = copy.deepcopy(vars(self))
54-
for k in ["SRCID", "DSTID", "type", "timestamp", "original_src_id1__", "original_dst_id2__"]:
55+
# for k in ["SRCID", "DSTID", "type", "timestamp", "original_src_id1__", "original_dst_id2__"]:
56+
for k in ["type", "original_src_id1__", "original_dst_id2__"]:
5557
attrs.pop(k)
5658
attrs.update(json.loads(attrs.pop("extra", '{}') or '{}'))
5759
return attrs

0 commit comments

Comments
 (0)