Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| from typing import Any | |
| from agent.agent import Agent | |
| from agent.agent_client import AgentClient, ChatOutcome, SkillsAgentClient | |
| from agent.champ_client import ChampAgentClient | |
| from agent.fake_client import FakeAgentClient | |
| from agent.frontier_client import ( | |
| GEMINI_CONSERVATIVE_TEMPERATURE, | |
| GEMINI_CREATIVE_TEMPERATURE, | |
| GeminiAgentClient, | |
| GptAgentClient, | |
| ) | |
| from classes.base_models import ChatRequest | |
| from providers.protocol import ChatProvider | |
| SessionId = str | |
| ConversationId = str | |
| ModelType = str | |
| class SessionRouter: | |
| """Per-(session, conversation) registry of AgentClient instances.""" | |
| def __init__( | |
| self, | |
| wiki_skills_agent: Agent, | |
| wiki_skills_agent_short: Agent, | |
| champ_provider: ChatProvider, | |
| champ_vector_store: Any, | |
| openai_provider: ChatProvider, | |
| openai_model_id: str, | |
| gemini_provider: ChatProvider, | |
| gemini_model_id: str, | |
| fake_provider: ChatProvider, | |
| ) -> None: | |
| self.session_clients: dict[tuple[SessionId, ConversationId], AgentClient] = {} | |
| self.wiki_skills_agent = wiki_skills_agent | |
| self.wiki_skills_agent_short = wiki_skills_agent_short | |
| self.champ_provider = champ_provider | |
| self.champ_vector_store = champ_vector_store | |
| self.openai_provider = openai_provider | |
| self.openai_model_id = openai_model_id | |
| self.gemini_provider = gemini_provider | |
| self.gemini_model_id = gemini_model_id | |
| self.fake_provider = fake_provider | |
| def delete_session(self, session_id: SessionId) -> None: | |
| """Drop every AgentClient (and its ConversationHistory) for this session.""" | |
| keys = [k for k in self.session_clients if k[0] == session_id] | |
| for k in keys: | |
| del self.session_clients[k] | |
| def send( | |
| self, payload: ChatRequest, documents: dict[str, str] | None = None | |
| ) -> ChatOutcome: | |
| # TODO: session/conversation ids should be created backend-side, not client-side. | |
| key = (payload.session_id, payload.conversation_id) | |
| agent_client = self.session_clients.get(key) | |
| if agent_client is None: | |
| agent_client = self._create_agent_client(payload.model_type) | |
| self.session_clients[key] = agent_client | |
| return agent_client.call( | |
| payload.human_message, lang=payload.lang, documents=documents | |
| ) | |
| def _create_agent_client(self, model_type: str) -> AgentClient: | |
| if model_type == "skills_wiki": | |
| return SkillsAgentClient(agent=self.wiki_skills_agent) | |
| if model_type == "skills_wiki_short": | |
| return SkillsAgentClient(agent=self.wiki_skills_agent_short) | |
| if model_type == "champ": | |
| return ChampAgentClient( | |
| provider=self.champ_provider, | |
| vector_store=self.champ_vector_store, | |
| ) | |
| if model_type == "openai": | |
| return GptAgentClient( | |
| provider=self.openai_provider, | |
| model_id=self.openai_model_id, | |
| ) | |
| if model_type == "google-conservative": | |
| return GeminiAgentClient( | |
| provider=self.gemini_provider, | |
| model_id=self.gemini_model_id, | |
| temperature=GEMINI_CONSERVATIVE_TEMPERATURE, | |
| ) | |
| if model_type == "google-creative": | |
| return GeminiAgentClient( | |
| provider=self.gemini_provider, | |
| model_id=self.gemini_model_id, | |
| temperature=GEMINI_CREATIVE_TEMPERATURE, | |
| ) | |
| if model_type == "fake": | |
| return FakeAgentClient(provider=self.fake_provider) | |
| raise RuntimeError(f"Unknown model_type: {model_type}") | |