NousResearch/atropos
1
0
Fork 0
mirror of https://github.com/NousResearch/atropos.git synced 2026-04-24 17:04:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Added flag for chat completion, switch in OpenAI client

Browse source
This commit is contained in:
Shannon Sands 2025-05-17 12:59:20 -07:00
parent 81534f7c8a
commit 59b08a1aa9
3 changed files with 45 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

15
atroposlib/envs/server_handling/openai_server.py
View file

@ -24,12 +24,19 @@ class OpenAIServer(APIServer):
)
super().__init__(config)
async def check_server_status_task(self):
async def check_server_status_task(self, chat_completion: bool = True):
while True:
try:
await self.openai.chat.completions.create(
model=self.config.model_name,
messages=[{"role": "user", "content": "hi"}],
if chat_completion:
await self.openai.chat.completions.create(
model=self.config.model_name,
messages=[{"role": "user", "content": "hi"}],
max_tokens=1,
)
else:
await self.openai.completions.create(
model=self.config.model_name,
prompt="hi",
max_tokens=1,
)
self.server_healthy = True