Spaces:

chrizefan
/

digichat

Runtime error

digichat / generators.py

Upload folder using huggingface_hub

fe52ef9 verified over 1 year ago

3.33 kB

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

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