1010import json
1111import logging
1212import pytest
13- from og_sdk .agent_sdk import AgentSDK
13+
14+ from og_sdk .kernel_sdk import KernelSDK
1415from og_agent .codellama_agent import CodellamaAgent
1516from og_proto .agent_server_pb2 import ProcessOptions , TaskResponse
1617import asyncio
2122logger = logging .getLogger (__name__ )
2223
2324
25+ @pytest .fixture
26+ def kernel_sdk ():
27+ endpoint = (
28+ "localhost:9527" # Replace with the actual endpoint of your test gRPC server
29+ )
30+ return KernelSDK (endpoint , "ZCeI9cYtOCyLISoi488BgZHeBkHWuFUH" )
31+
32+
2433class PayloadStream :
2534
2635 def __init__ (self , payload ):
@@ -49,24 +58,157 @@ def done(self):
4958
5059class CodellamaMockClient :
5160
52- def __init__ (self , payload ):
53- self .payload = payload
61+ def __init__ (self , payloads ):
62+ self .payloads = payloads
63+ self .index = 0
5464
5565 async def prompt (self , question , chat_history = []):
56- async for line in PayloadStream (self .payload ):
66+ if self .index >= len (self .payloads ):
67+ raise StopAsyncIteration
68+ self .index += 1
69+ payload = self .payloads [self .index - 1 ]
70+ async for line in PayloadStream (payload ):
5771 yield line
5872
5973
60- @pytest_asyncio .fixture
61- async def agent_sdk ():
62- sdk = AgentSDK (api_base , api_key )
63- sdk .connect ()
64- yield sdk
65- await sdk .close ()
74+ @pytest .mark .asyncio
75+ async def test_codellama_agent_execute_bash_code (kernel_sdk ):
76+ kernel_sdk .connect ()
77+ sentence1 = {
78+ "explanation" : "print a hello world using python" ,
79+ "action" : "execute_bash_code" ,
80+ "action_input" : "echo 'hello world'" ,
81+ "saved_filenames" : [],
82+ "language" : "python" ,
83+ "is_final_answer" : False ,
84+ }
85+ sentence2 = {
86+ "explanation" : "the output matchs the goal" ,
87+ "action" : "no_action" ,
88+ "action_input" : "" ,
89+ "saved_filenames" : [],
90+ "language" : "en" ,
91+ "is_final_answer" : False ,
92+ }
93+ client = CodellamaMockClient ([json .dumps (sentence1 ), json .dumps (sentence2 )])
94+ agent = CodellamaAgent (client , kernel_sdk )
95+ task_opt = ProcessOptions (
96+ streaming = True ,
97+ llm_name = "codellama" ,
98+ input_token_limit = 100000 ,
99+ output_token_limit = 100000 ,
100+ timeout = 5 ,
101+ )
102+ queue = asyncio .Queue ()
103+ await agent .arun ("write a hello world in bash" , queue , MockContext (), task_opt )
104+ responses = []
105+ while True :
106+ try :
107+ response = await queue .get ()
108+ if not response :
109+ break
110+ responses .append (response )
111+ except asyncio .QueueEmpty :
112+ break
113+ logger .info (responses )
114+ console_output = list (
115+ filter (
116+ lambda x : x .response_type == TaskResponse .OnStepActionStreamStdout ,
117+ responses ,
118+ )
119+ )
120+ assert len (console_output ) == 1 , "bad console output count"
121+ assert console_output [0 ].console_stdout == "hello world\n " , "bad console output"
122+
123+
124+ @pytest .mark .asyncio
125+ async def test_codellama_agent_execute_python_code (kernel_sdk ):
126+ kernel_sdk .connect ()
127+ sentence1 = {
128+ "explanation" : "print a hello world using python" ,
129+ "action" : "execute_python_code" ,
130+ "action_input" : "print('hello world')" ,
131+ "saved_filenames" : [],
132+ "language" : "python" ,
133+ "is_final_answer" : False ,
134+ }
135+ sentence2 = {
136+ "explanation" : "the output matchs the goal" ,
137+ "action" : "no_action" ,
138+ "action_input" : "" ,
139+ "saved_filenames" : [],
140+ "language" : "en" ,
141+ "is_final_answer" : False ,
142+ }
143+ client = CodellamaMockClient ([json .dumps (sentence1 ), json .dumps (sentence2 )])
144+ agent = CodellamaAgent (client , kernel_sdk )
145+ task_opt = ProcessOptions (
146+ streaming = True ,
147+ llm_name = "codellama" ,
148+ input_token_limit = 100000 ,
149+ output_token_limit = 100000 ,
150+ timeout = 5 ,
151+ )
152+ queue = asyncio .Queue ()
153+ await agent .arun ("write a hello world in python" , queue , MockContext (), task_opt )
154+ responses = []
155+ while True :
156+ try :
157+ response = await queue .get ()
158+ if not response :
159+ break
160+ responses .append (response )
161+ except asyncio .QueueEmpty :
162+ break
163+ logger .info (responses )
164+ console_output = list (
165+ filter (
166+ lambda x : x .response_type == TaskResponse .OnStepActionStreamStdout ,
167+ responses ,
168+ )
169+ )
170+ assert len (console_output ) == 1 , "bad console output count"
171+ assert console_output [0 ].console_stdout == "hello world\n " , "bad console output"
172+
173+
174+ @pytest .mark .asyncio
175+ async def test_codellama_agent_show_demo_code (kernel_sdk ):
176+ sentence = {
177+ "explanation" : "Hello, how can I help you?" ,
178+ "action" : "show_demo_code" ,
179+ "action_input" : "echo 'hello world'" ,
180+ "saved_filenames" : [],
181+ "language" : "shell" ,
182+ "is_final_answer" : True ,
183+ }
184+ client = CodellamaMockClient ([json .dumps (sentence )])
185+ agent = CodellamaAgent (client , kernel_sdk )
186+ task_opt = ProcessOptions (
187+ streaming = True ,
188+ llm_name = "codellama" ,
189+ input_token_limit = 100000 ,
190+ output_token_limit = 100000 ,
191+ timeout = 5 ,
192+ )
193+ queue = asyncio .Queue ()
194+ await agent .arun ("hello" , queue , MockContext (), task_opt )
195+ responses = []
196+ while True :
197+ try :
198+ response = await queue .get ()
199+ if not response :
200+ break
201+ responses .append (response )
202+ except asyncio .QueueEmpty :
203+ break
204+ logger .info (responses )
205+ assert (
206+ responses [- 1 ].response_type == TaskResponse .OnFinalAnswer
207+ ), "bad response type"
66208
67209
68210@pytest .mark .asyncio
69- async def test_codellama_agent_smoke_test (agent_sdk ):
211+ async def test_codellama_agent_smoke_test (kernel_sdk ):
70212 sentence = {
71213 "explanation" : "Hello, how can I help you?" ,
72214 "action" : "no_action" ,
@@ -75,8 +217,8 @@ async def test_codellama_agent_smoke_test(agent_sdk):
75217 "language" : "en" ,
76218 "is_final_answer" : True ,
77219 }
78- client = CodellamaMockClient (json .dumps (sentence ))
79- agent = CodellamaAgent (client , agent_sdk )
220+ client = CodellamaMockClient ([ json .dumps (sentence )] )
221+ agent = CodellamaAgent (client , kernel_sdk )
80222 task_opt = ProcessOptions (
81223 streaming = True ,
82224 llm_name = "codellama" ,
0 commit comments