generators.py

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}}]}