File size: 3,326 Bytes
fe52ef9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | import os
from typing import List, Dict, Any, Generator
import importlib
import threading
import queue
from dotenv import load_dotenv
from client_utils import get_client, get_model, get_config, engine_map
load_dotenv()
def get_agent(agent_type: str, *args, **kwargs):
"""Retrieve an agent configuration based on the agent type from engine_map."""
try:
if agent_type.lower() in engine_map:
return engine_map[agent_type.lower()]["config"]
raise ValueError(f"Unsupported agent type: {agent_type}")
except KeyError:
raise ValueError(f"Unsupported agent type: {agent_type}")
def chat_completion(messages: List[Dict[str, Any]], engine: str = "grok-2-latest", temperature: float = 0.5) -> str:
# Dynamic import to avoid circular import
from agents import BaseAgent, OrchestratorAgent
# Get the appropriate config
config = get_config(engine)
# Override temperature if specified
config.setup_model_config(engine, temperature)
# Use orchestrator if the config has tools available
if len(config.get_tools()) > 0:
agent = OrchestratorAgent(config)
else:
agent = BaseAgent(config)
result_queue = queue.Queue()
def run_process():
try:
result = agent.process(messages)
result_queue.put(result)
except Exception as e:
result_queue.put(e)
t = threading.Thread(target=run_process)
t.start()
t.join()
result = result_queue.get()
if isinstance(result, Exception):
raise result
return result
def chat_completion_stream(messages: List[Dict[str, Any]], engine: str = "grok-2-latest", temperature: float = 0.5) -> Generator[Any, None, None]:
"""Stream chat completions based on the selected engine or agent."""
try:
# Dynamic import to avoid circular import
from agents import BaseAgent, OrchestratorAgent
# Get the appropriate config
config = get_config(engine)
# Override temperature if specified
config.setup_model_config(engine, temperature)
# Use orchestrator if the config has tools available
if len(config.get_tools()) > 0:
agent = OrchestratorAgent(config)
stream_fn = agent.process_stream
else:
agent = BaseAgent(config)
stream_fn = agent.process_stream
q = queue.Queue()
sentinel = object()
def run_stream():
try:
for chunk in stream_fn(messages):
q.put(chunk)
except Exception as e:
q.put(e)
finally:
q.put(sentinel)
t = threading.Thread(target=run_stream)
t.start()
while True:
item = q.get()
if item is sentinel:
break
if isinstance(item, Exception):
raise item
yield item
except Exception as e:
# In case of any error, yield a special error chunk that our handler can process
print(f"Error in chat_completion_stream: {e}")
error_msg = f"Error: {str(e)}"
# Create a minimal mock chunk with the error message
yield {"choices": [{"delta": {"content": error_msg}}]}
|