From a9a7b25103d16e1b4d99f127397842ae2a7eac58 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 18 Jul 2026 12:01:34 +0800 Subject: [PATCH 01/25] Refactor/chat streamer (#1) * refactor chat streaming * update `_is_token_limit_error` logic * test: add chat provider unittests * move `_is_token_limit_error` to module level * remove unnecessary TYPE_CHECKING * add EOF * update prompt function --- api/chat.py | 279 ++++++++++++++++++++++ api/simple_chat.py | 515 ++++++---------------------------------- tests/unit/test_chat.py | 41 ++++ 3 files changed, 389 insertions(+), 446 deletions(-) create mode 100644 api/chat.py create mode 100644 tests/unit/test_chat.py diff --git a/api/chat.py b/api/chat.py new file mode 100644 index 000000000..60b05aba8 --- /dev/null +++ b/api/chat.py @@ -0,0 +1,279 @@ +import logging +from abc import abstractmethod, ABC + +from typing import TYPE_CHECKING + +from collections.abc import AsyncIterator + +from adalflow.core.types import ModelType +from api.config import ( + OPENROUTER_API_KEY, + OPENAI_API_KEY, + AWS_ACCESS_KEY_ID, + AWS_SECRET_ACCESS_KEY, +) + +if TYPE_CHECKING: + from ollama import ChatResponse + from openai.types.chat import ChatCompletionChunk + from openai import AsyncStream + +MODEL_CFG = dict[str, str | int | float] + +# Configure logging +from api.logging_config import setup_logging + +setup_logging() +logger = logging.getLogger(__name__) + + +class ChatStreamer(ABC): + _registry: dict[str, type["ChatStreamer"]] = {} + provider: str + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + if provider := getattr(cls, "provider", None): + ChatStreamer._registry[provider] = cls + + @classmethod + def create(cls, *, provider: str, model: str | None = None, model_config: MODEL_CFG) -> "ChatStreamer": + model = model or model_config.get("model") + logger.info("Using %s with model: %s", provider, model) + registered = ChatStreamer._registry.get(provider, None) + if registered: + return registered(model=model, model_config=model_config) + raise RuntimeError(f"Provider {provider} not registered") + + @abstractmethod + def respond_stream(self, prompt: str) -> AsyncIterator[str]: + raise NotImplementedError(f"{type(self).__name__} does not implement `respond_stream`") + + +class OllamaChatStreamer(ChatStreamer): + provider = "ollama" + + def __init__(self, *, model: str, model_config: MODEL_CFG): + from adalflow.components.model_client.ollama_client import OllamaClient + + self.client = OllamaClient() + self.model_kwargs = { + "model": model, + "stream": True, + "options": { + "temperature": model_config["temperature"], + "top_p": model_config["top_p"], + "num_ctx": model_config["num_ctx"] + } + } + + logger.debug(f"Prompting Ollama with kwargs: {self.model_kwargs}") + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt + " /no_think", # todo I think this could be added into model kwargs? + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM, + ) + + response: "AsyncIterator[ChatResponse]" = await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ) + + async for chunk in response: + if not hasattr(chunk, "message"): + raise RuntimeError( + "`message` field not found in response. Wrong ollama-python version probably.", + ) + text = chunk.message.content + if text: + text = text.replace('', '').replace('', '') + yield text + + +class OpenRouterChatStreamer(ChatStreamer): + provider = "openrouter" + + def __init__(self, *, model: str, model_config: MODEL_CFG): + if not OPENROUTER_API_KEY: + logger.warning("OPENROUTER_API_KEY not configured, but continuing with request") + # We'll let the OpenRouterClient handle this and return a friendly error message + from api.openrouter_client import OpenRouterClient + + self.client = OpenRouterClient() + self.model_kwargs = { + "model": model, + "stream": True, + "temperature": model_config["temperature"] + } + if "top_k" in model_config: + self.model_kwargs["top_k"] = model_config["top_k"] + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt, + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM, + ) + async for chunk in await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ): + yield chunk + + +class OpenAIChatStreamer(ChatStreamer): + provider = "openai" + + def __init__(self, *, model: str, model_config: MODEL_CFG): + if not OPENAI_API_KEY: + logger.warning("OPENAI_API_KEY not configured, but continuing with request") + # We'll let the OpenAIClient handle this and return an error message + from api.openai_client import OpenAIClient + + self.client = OpenAIClient() + self.model_kwargs = { + "model": model, + "stream": True, + "temperature": model_config["temperature"] + } + # Only add top_p if it exists in the model config + if "top_p" in model_config: + self.model_kwargs["top_p"] = model_config["top_p"] + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt, + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM + ) + response: "AsyncStream[ChatCompletionChunk]" = await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ) + + async for chunk in response: + if ( + chunk.choices and + chunk.choices[0].delta is not None and + chunk.choices[0].delta.content is not None + ): + yield chunk.choices[0].delta.content + + +class AzureChatStreamer(ChatStreamer): + provider = "azure" + + def __init__(self, *, model: str, model_config: MODEL_CFG): + from api.azureai_client import AzureAIClient + + self.client = AzureAIClient() + self.model_kwargs = { + "model": model, + "stream": True, + "temperature": model_config["temperature"], + "top_p": model_config["top_p"] + } + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt, + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM + ) + response: "AsyncStream[ChatCompletionChunk]" = await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ) + + async for chunk in response: + if ( + chunk.choices and + chunk.choices[0].delta is not None and + chunk.choices[0].delta.content is not None + ): + yield chunk.choices[0].delta.content + + +class BedrockChatStreamer(ChatStreamer): + provider = "bedrock" + + def __init__(self, *, model: str, model_config: MODEL_CFG): + if not AWS_ACCESS_KEY_ID or not AWS_SECRET_ACCESS_KEY: + logger.warning("AWS_ACCESS_KEY_ID or AWS_SECRET_ACCESS_KEY not configured, but continuing with request") + # We'll let the BedrockClient handle this and return an error message + from api.bedrock_client import BedrockClient + + self.client = BedrockClient() + self.model_kwargs = { + "model": model, + "temperature": model_config["temperature"], + "top_p": model_config["top_p"] + } + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt, + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM, + ) + response = await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ) + if not isinstance(response, str): + response = str(response) + yield response + + +class DashScopeChatStreamer(ChatStreamer): + provider = "dashscope" + + def __init__(self, *, model: str, model_config: MODEL_CFG): + from api.dashscope_client import DashscopeClient + + self.client = DashscopeClient() + self.model_kwargs = { + "model": model, + "stream": True, + "temperature": model_config["temperature"], + "top_p": model_config["top_p"], + } + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt, + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM, + ) + response = await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ) + async for text in response: + if text: + yield text + + +class GoogleGenerativeChatStreamer(ChatStreamer): + provider = "google" + + def __init__(self, *, model: str, model_config: MODEL_CFG): + import google.generativeai as genai + from google.generativeai.types import GenerationConfig + + self.client = genai.GenerativeModel( + model_name=model, + generation_config=GenerationConfig( + temperature=model_config["temperature"], + top_p=model_config["top_p"], + top_k=model_config["top_k"], + ) + ) + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + response = await self.client.generate_content_async(prompt, stream=True) + async for chunk in response: + if hasattr(chunk, "text"): + yield chunk.text diff --git a/api/simple_chat.py b/api/simple_chat.py index e93199397..cecb146e1 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -1,24 +1,15 @@ import logging -import os -from collections.abc import AsyncIterator -from typing import List, Optional, TYPE_CHECKING +from typing import List, Optional, Callable from urllib.parse import unquote -import google.generativeai as genai -from adalflow.components.model_client.ollama_client import OllamaClient -from adalflow.core.types import ModelType from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field -from api.config import get_model_config, configs, OPENROUTER_API_KEY, OPENAI_API_KEY, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY +from api.chat import ChatStreamer +from api.config import get_model_config, configs from api.data_pipeline import count_tokens, get_file_content -from api.openai_client import OpenAIClient -from api.openrouter_client import OpenRouterClient -from api.bedrock_client import BedrockClient -from api.azureai_client import AzureAIClient -from api.dashscope_client import DashscopeClient from api.rag import RAG from api.prompts import ( DEEP_RESEARCH_FIRST_ITERATION_PROMPT, @@ -30,9 +21,6 @@ # Configure logging from api.logging_config import setup_logging -if TYPE_CHECKING: - from ollama import ChatResponse - setup_logging() logger = logging.getLogger(__name__) @@ -77,6 +65,14 @@ class ChatCompletionRequest(BaseModel): included_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to include exclusively") included_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to include exclusively") + +def _is_token_limit_error(exc: Exception) -> bool: + error_message = str(exc).lower() + return any( + k in error_message for k in ("maximum context length", "token limit", "too many tokens") + ) + + @app.post("/chat/completions/stream") async def chat_completions_stream(request: ChatCompletionRequest): """Stream a chat completion response directly using Google Generative AI""" @@ -308,453 +304,79 @@ async def chat_completions_stream(request: ChatCompletionRequest): if not isinstance(turn_id, int) and hasattr(turn, 'user_query') and hasattr(turn, 'assistant_response'): conversation_history += f"\n{turn.user_query.query_str}\n{turn.assistant_response.response_str}\n\n" - # Create the prompt with context - prompt = f"/no_think {system_prompt}\n\n" + def prompt() -> str: + # Create the prompt with context + prompt = f"/no_think {system_prompt}\n\n" + if conversation_history: + prompt += f"\n{conversation_history}\n\n" - if conversation_history: - prompt += f"\n{conversation_history}\n\n" + # Check if filePath is provided and fetch file content if it exists + if request.filePath and file_content: + # Add file content to the prompt after conversation history + prompt += f"\n{file_content}\n\n\n" - # Check if filePath is provided and fetch file content if it exists - if file_content: - # Add file content to the prompt after conversation history - prompt += f"\n{file_content}\n\n\n" + # Only include context if it's not empty + if context_text.strip(): + context_prompt = f"\n{context_text}\n\n\n" + else: + # Add a note that we're skipping RAG due to size constraints or because it's the isolated API + logger.info("No context available from RAG") + context_prompt = "Answering without retrieval augmentation.\n\n" - # Only include context if it's not empty - CONTEXT_START = "" - CONTEXT_END = "" - if context_text.strip(): - prompt += f"{CONTEXT_START}\n{context_text}\n{CONTEXT_END}\n\n" - else: - # Add a note that we're skipping RAG due to size constraints or because it's the isolated API - logger.info("No context available from RAG") - prompt += "Answering without retrieval augmentation.\n\n" + prompt += f"{context_prompt}\n{query}\n\n\nAssistant: " + return prompt - prompt += f"\n{query}\n\n\nAssistant: " + def simplified_prompt() -> str: + prompt = f"/no_think {system_prompt}\n\n" + if conversation_history: + prompt += f"\n{conversation_history}\n\n" - model_config = get_model_config(request.provider, request.model)["model_kwargs"] + # Include file content in the fallback prompt if it was retrieved + if request.filePath and file_content: + prompt += f"\n{file_content}\n\n\n" - if request.provider == "ollama": - prompt += " /no_think" - - model = OllamaClient() - model_kwargs = { - "model": model_config["model"], - "stream": True, - "options": { - "temperature": model_config["temperature"], - "top_p": model_config["top_p"], - "num_ctx": model_config["num_ctx"] - } - } - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "openrouter": - logger.info(f"Using OpenRouter with model: {request.model}") - - # Check if OpenRouter API key is set - if not OPENROUTER_API_KEY: - logger.warning("OPENROUTER_API_KEY not configured, but continuing with request") - # We'll let the OpenRouterClient handle this and return a friendly error message - - model = OpenRouterClient() - model_kwargs = { - "model": request.model, - "stream": True, - "temperature": model_config["temperature"] - } - # Only add top_p if it exists in the model config - if "top_p" in model_config: - model_kwargs["top_p"] = model_config["top_p"] - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "openai": - logger.info(f"Using Openai protocol with model: {request.model}") - - # Check if an API key is set for Openai - if not OPENAI_API_KEY: - logger.warning("OPENAI_API_KEY not configured, but continuing with request") - # We'll let the OpenAIClient handle this and return an error message - - # Initialize Openai client - model = OpenAIClient() - model_kwargs = { - "model": request.model, - "stream": True, - "temperature": model_config["temperature"] - } - # Only add top_p if it exists in the model config - if "top_p" in model_config: - model_kwargs["top_p"] = model_config["top_p"] - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "bedrock": - logger.info(f"Using AWS Bedrock with model: {request.model}") - - # Check if AWS credentials are set - if not AWS_ACCESS_KEY_ID or not AWS_SECRET_ACCESS_KEY: - logger.warning("AWS_ACCESS_KEY_ID or AWS_SECRET_ACCESS_KEY not configured, but continuing with request") - # We'll let the BedrockClient handle this and return an error message - - # Initialize Bedrock client - model = BedrockClient() - model_kwargs = { - "model": request.model, - "temperature": model_config["temperature"], - "top_p": model_config["top_p"] - } - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "azure": - logger.info(f"Using Azure AI with model: {request.model}") - - # Initialize Azure AI client - model = AzureAIClient() - model_kwargs = { - "model": request.model, - "stream": True, - "temperature": model_config["temperature"], - "top_p": model_config["top_p"] - } - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "dashscope": - logger.info(f"Using Dashscope with model: {request.model}") - - model = DashscopeClient() - model_kwargs = { - "model": request.model, - "stream": True, - "temperature": model_config["temperature"], - "top_p": model_config["top_p"], - } - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM, - ) - else: - # Initialize Google Generative AI model (default provider) - model = genai.GenerativeModel( - model_name=model_config["model"], - generation_config={ - "temperature": model_config["temperature"], - "top_p": model_config["top_p"], - "top_k": model_config["top_k"], - }, + prompt += ( + f"Answering without retrieval augmentation due to input size constraints.\n\n" + f"\n{query}\n\n\nAssistant: " ) + return prompt - # Create a streaming response - async def response_stream(): + async def stream_and_fallback( + streamer: ChatStreamer, + prompt_func: Callable[[], str], + simplified_prompt_func: Callable[[], str], + ): try: - if request.provider == "ollama": - # Get the response and handle it properly using the previously created api_kwargs - response: "AsyncIterator[ChatResponse]" = await model.acall( - api_kwargs=api_kwargs, - model_type=ModelType.LLM, - ) - # Handle streaming response from Ollama - async for chunk in response: - if not hasattr(chunk, "message"): - raise RuntimeError( - "`message` field not found in response. Wrong ollama-python version probably.", - ) - text = chunk.message.content - if text: - text = text.replace('', '').replace('', '') - yield text - elif request.provider == "openrouter": - try: - # Get the response and handle it properly using the previously created api_kwargs - logger.info("Making OpenRouter API call") - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - # Handle streaming response from OpenRouter - async for chunk in response: - yield chunk - except Exception as e_openrouter: - logger.error(f"Error with OpenRouter API: {str(e_openrouter)}") - yield f"\nError with OpenRouter API: {str(e_openrouter)}\n\nPlease check that you have set the OPENROUTER_API_KEY environment variable with a valid API key." - elif request.provider == "openai": - try: - # Get the response and handle it properly using the previously created api_kwargs - logger.info("Making Openai API call") - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - # Handle streaming response from Openai - async for chunk in response: - choices = getattr(chunk, "choices", []) - if len(choices) > 0: - delta = getattr(choices[0], "delta", None) - if delta is not None: - text = getattr(delta, "content", None) - if text is not None: - yield text - except Exception as e_openai: - logger.error(f"Error with Openai API: {str(e_openai)}") - yield f"\nError with Openai API: {str(e_openai)}\n\nPlease check that you have set the OPENAI_API_KEY environment variable with a valid API key." - elif request.provider == "bedrock": - try: - # Get the response and handle it properly using the previously created api_kwargs - logger.info("Making AWS Bedrock API call") - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - # Handle response from Bedrock (not streaming yet) - if isinstance(response, str): - yield response - else: - # Try to extract text from the response - yield str(response) - except Exception as e_bedrock: - logger.error(f"Error with AWS Bedrock API: {str(e_bedrock)}") - yield f"\nError with AWS Bedrock API: {str(e_bedrock)}\n\nPlease check that you have set the AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables with valid credentials." - elif request.provider == "azure": - try: - # Get the response and handle it properly using the previously created api_kwargs - logger.info("Making Azure AI API call") - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - # Handle streaming response from Azure AI - async for chunk in response: - choices = getattr(chunk, "choices", []) - if len(choices) > 0: - delta = getattr(choices[0], "delta", None) - if delta is not None: - text = getattr(delta, "content", None) - if text is not None: - yield text - except Exception as e_azure: - logger.error(f"Error with Azure AI API: {str(e_azure)}") - yield f"\nError with Azure AI API: {str(e_azure)}\n\nPlease check that you have set the AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, and AZURE_OPENAI_VERSION environment variables with valid values." - elif request.provider == "dashscope": - try: - logger.info("Making Dashscope API call") - response = await model.acall( - api_kwargs=api_kwargs, model_type=ModelType.LLM - ) - # DashscopeClient.acall with stream=True returns an async - # generator of text chunks - async for text in response: - if text: - yield text - except Exception as e_dashscope: - logger.error(f"Error with Dashscope API: {str(e_dashscope)}") - yield ( - f"\nError with Dashscope API: {str(e_dashscope)}\n\n" - "Please check that you have set the DASHSCOPE_API_KEY (and optionally " - "DASHSCOPE_WORKSPACE_ID) environment variables with valid values." - ) - else: - # Google Generative AI (default provider) - response = model.generate_content(prompt, stream=True) - for chunk in response: - if hasattr(chunk, "text"): - yield chunk.text - - except Exception as e_outer: - logger.error(f"Error in streaming response: {str(e_outer)}") - error_message = str(e_outer) - - # Check for token limit errors - if "maximum context length" in error_message or "token limit" in error_message or "too many tokens" in error_message: - # If we hit a token limit error, try again without context + async for chunk in streamer.respond_stream(prompt_func()): + yield chunk + except Exception as e: + if _is_token_limit_error(e): logger.warning("Token limit exceeded, retrying without context") try: - # Create a simplified prompt without context - simplified_prompt = f"/no_think {system_prompt}\n\n" - if conversation_history: - simplified_prompt += f"\n{conversation_history}\n\n" - - # Include file content in the fallback prompt if it was retrieved - if request.filePath and file_content: - simplified_prompt += f"\n{file_content}\n\n\n" - - simplified_prompt += "Answering without retrieval augmentation due to input size constraints.\n\n" - simplified_prompt += f"\n{query}\n\n\nAssistant: " - - if request.provider == "ollama": - simplified_prompt += " /no_think" - - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - fallback_response: "AsyncIterator[ChatResponse]" = await model.acall( - api_kwargs=fallback_api_kwargs, - model_type=ModelType.LLM, - ) - - # Handle streaming fallback_response from Ollama - async for chunk in fallback_response: - if not hasattr(chunk, "message"): - raise RuntimeError( - "`message` field not found in response. Wrong ollama-python version probably.", - ) - text = chunk.message.content - if text: - text = text.replace('', '').replace('', '') - yield text - elif request.provider == "openrouter": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - logger.info("Making fallback OpenRouter API call") - fallback_response = await model.acall(api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM) - - # Handle streaming fallback_response from OpenRouter - async for chunk in fallback_response: - yield chunk - except Exception as e_fallback: - logger.error(f"Error with OpenRouter API fallback: {str(e_fallback)}") - yield f"\nError with OpenRouter API fallback: {str(e_fallback)}\n\nPlease check that you have set the OPENROUTER_API_KEY environment variable with a valid API key." - elif request.provider == "openai": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - logger.info("Making fallback Openai API call") - fallback_response = await model.acall(api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM) - - # Handle streaming fallback_response from Openai - async for chunk in fallback_response: - text = chunk if isinstance(chunk, str) else getattr(chunk, 'text', str(chunk)) - yield text - except Exception as e_fallback: - logger.error(f"Error with Openai API fallback: {str(e_fallback)}") - yield f"\nError with Openai API fallback: {str(e_fallback)}\n\nPlease check that you have set the OPENAI_API_KEY environment variable with a valid API key." - elif request.provider == "bedrock": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - logger.info("Making fallback AWS Bedrock API call") - fallback_response = await model.acall(api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM) - - # Handle response from Bedrock - if isinstance(fallback_response, str): - yield fallback_response - else: - # Try to extract text from the response - yield str(fallback_response) - except Exception as e_fallback: - logger.error(f"Error with AWS Bedrock API fallback: {str(e_fallback)}") - yield f"\nError with AWS Bedrock API fallback: {str(e_fallback)}\n\nPlease check that you have set the AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables with valid credentials." - elif request.provider == "azure": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - logger.info("Making fallback Azure AI API call") - fallback_response = await model.acall(api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM) - - # Handle streaming fallback response from Azure AI - async for chunk in fallback_response: - choices = getattr(chunk, "choices", []) - if len(choices) > 0: - delta = getattr(choices[0], "delta", None) - if delta is not None: - text = getattr(delta, "content", None) - if text is not None: - yield text - except Exception as e_fallback: - logger.error(f"Error with Azure AI API fallback: {str(e_fallback)}") - yield f"\nError with Azure AI API fallback: {str(e_fallback)}\n\nPlease check that you have set the AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, and AZURE_OPENAI_VERSION environment variables with valid values." - elif request.provider == "dashscope": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM, - ) - - logger.info("Making fallback Dashscope API call") - fallback_response = await model.acall( - api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM - ) - - # DashscopeClient.acall (stream=True) returns an async - # generator of text chunks - async for text in fallback_response: - if text: - yield text - except Exception as e_fallback: - logger.error( - f"Error with Dashscope API fallback: {str(e_fallback)}" - ) - yield ( - f"\nError with Dashscope API fallback: {str(e_fallback)}\n\n" - "Please check that you have set the DASHSCOPE_API_KEY (and optionally " - "DASHSCOPE_WORKSPACE_ID) environment variables with valid values." - ) - else: - # Google Generative AI fallback (default provider) - model_config = get_model_config(request.provider, request.model) - fallback_model = genai.GenerativeModel( - model_name=model_config["model_kwargs"]["model"], - generation_config={ - "temperature": model_config["model_kwargs"].get("temperature", 0.7), - "top_p": model_config["model_kwargs"].get("top_p", 0.8), - "top_k": model_config["model_kwargs"].get("top_k", 40), - }, - ) - - fallback_response = fallback_model.generate_content( - simplified_prompt, stream=True - ) - for chunk in fallback_response: - if hasattr(chunk, "text"): - yield chunk.text + async for chunk in streamer.respond_stream(simplified_prompt_func()): + yield chunk except Exception as e2: logger.error(f"Error in fallback streaming response: {str(e2)}") - yield f"\nI apologize, but your request is too large for me to process. Please try a shorter query or break it into smaller parts." + yield ( + f"\nI apologize, but your request is too large for me to process. " + f"Please try a shorter query or break it into smaller parts." + ) else: - # For other errors, return the error message - yield f"\nError: {error_message}" + yield f"\nError: {str(e)}" + + model_config = get_model_config(request.provider, request.model)["model_kwargs"] + chat_streamer = ChatStreamer.create( + provider=request.provider, + model=request.model, + model_config=model_config, + ) # Return streaming response - return StreamingResponse(response_stream(), media_type="text/event-stream") + return StreamingResponse(stream_and_fallback( + streamer=chat_streamer, + prompt_func=prompt, + simplified_prompt_func=simplified_prompt, + ), media_type="text/event-stream") except HTTPException: raise @@ -763,6 +385,7 @@ async def response_stream(): logger.error(error_msg) raise HTTPException(status_code=500, detail=error_msg) + @app.get("/") async def root(): """Root endpoint to check if the API is running""" diff --git a/tests/unit/test_chat.py b/tests/unit/test_chat.py new file mode 100644 index 000000000..181791626 --- /dev/null +++ b/tests/unit/test_chat.py @@ -0,0 +1,41 @@ +import pytest +from api.chat import ( + ChatStreamer, + OllamaChatStreamer, + OpenRouterChatStreamer, + OpenAIChatStreamer, + AzureChatStreamer, + BedrockChatStreamer, + DashScopeChatStreamer, + GoogleGenerativeChatStreamer, +) + +@pytest.mark.parametrize("provider, expected", [ + ("ollama", OllamaChatStreamer), + ("openrouter", OpenRouterChatStreamer), + ("openai", OpenAIChatStreamer), + ("azure", AzureChatStreamer), + ("bedrock", BedrockChatStreamer), + ("dashscope", DashScopeChatStreamer), + ("google", GoogleGenerativeChatStreamer), +]) +def test_every_provider_is_registered(provider, expected): + assert ChatStreamer._registry[provider] is expected + +@pytest.mark.parametrize("provider, expected", [ + ("ollama", OllamaChatStreamer), + ("openrouter", OpenRouterChatStreamer), + ("openai", OpenAIChatStreamer), + ("azure", AzureChatStreamer), + ("bedrock", BedrockChatStreamer), + ("dashscope", DashScopeChatStreamer), + ("google", GoogleGenerativeChatStreamer), +]) +def test_create_returns_correct_subclass(monkeypatch, provider, expected): + monkeypatch.setattr(expected, "__init__", lambda self, **kw: None) + s = ChatStreamer.create(provider=provider, model="m", model_config={"model": "m"}) + assert isinstance(s, expected) + +def test_create_unknown_provider_raises(): + with pytest.raises(RuntimeError, match="not registered"): + ChatStreamer.create(provider="nope", model=None, model_config={}) From 149c6eba2624ca6e10c79ac1360a10f929003227 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 18 Jul 2026 12:08:44 +0800 Subject: [PATCH 02/25] update `input_too_large` threshold (#2) --- api/simple_chat.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/api/simple_chat.py b/api/simple_chat.py index cecb146e1..fb7fe4d07 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -10,7 +10,7 @@ from api.chat import ChatStreamer from api.config import get_model_config, configs from api.data_pipeline import count_tokens, get_file_content -from api.rag import RAG +from api.rag import RAG, MAX_INPUT_TOKENS from api.prompts import ( DEEP_RESEARCH_FIRST_ITERATION_PROMPT, DEEP_RESEARCH_FINAL_ITERATION_PROMPT, @@ -84,8 +84,8 @@ async def chat_completions_stream(request: ChatCompletionRequest): if hasattr(last_message, 'content') and last_message.content: tokens = count_tokens(last_message.content, request.provider == "ollama") logger.info(f"Request size: {tokens} tokens") - if tokens > 8000: - logger.warning(f"Request exceeds recommended token limit ({tokens} > 7500)") + if tokens > MAX_INPUT_TOKENS: + logger.warning(f"Request exceeds recommended token limit ({tokens} > {MAX_INPUT_TOKENS})") input_too_large = True # Create a new RAG instance for this request From 6bdc4e8ddb53a0ca705550b5eda17395bab34d83 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 18 Jul 2026 21:05:23 +0800 Subject: [PATCH 03/25] patch: patch adalflow OllamaClient to support batch embedding api. (#3) related to https://github.com/SylphAI-Inc/AdalFlow/pull/496 --- api/config.py | 3 +- api/data_pipeline.py | 14 ++--- api/ollama_client.py | 131 +++++++++++++++++++++++++++++++++++++++++++ api/rag.py | 19 +------ 4 files changed, 138 insertions(+), 29 deletions(-) create mode 100644 api/ollama_client.py diff --git a/api/config.py b/api/config.py index 1d6a5874c..4cafcdbc7 100644 --- a/api/config.py +++ b/api/config.py @@ -14,7 +14,8 @@ from api.google_embedder_client import GoogleEmbedderClient from api.azureai_client import AzureAIClient from api.dashscope_client import DashscopeClient -from adalflow import GoogleGenAIClient, OllamaClient +from api.ollama_client import OllamaClient +from adalflow import GoogleGenAIClient # Get API keys from environment variables OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') diff --git a/api/data_pipeline.py b/api/data_pipeline.py index f98068651..1d88ef9f7 100644 --- a/api/data_pipeline.py +++ b/api/data_pipeline.py @@ -415,16 +415,10 @@ def prepare_data_pipeline(embedder_type: str = None, is_ollama_embedder: bool = embedder = get_embedder(embedder_type=embedder_type) - # Choose appropriate processor based on embedder type - if embedder_type == 'ollama': - # Use Ollama document processor for single-document processing - embedder_transformer = OllamaDocumentProcessor(embedder=embedder) - else: - # Use batch processing for OpenAI and Google embedders - batch_size = embedder_config.get("batch_size", 500) - embedder_transformer = ToEmbeddings( - embedder=embedder, batch_size=batch_size - ) + batch_size = embedder_config.get("batch_size", 500) + embedder_transformer = ToEmbeddings( + embedder=embedder, batch_size=batch_size + ) data_transformer = adal.Sequential( splitter, embedder_transformer diff --git a/api/ollama_client.py b/api/ollama_client.py new file mode 100644 index 000000000..2418e313d --- /dev/null +++ b/api/ollama_client.py @@ -0,0 +1,131 @@ +# This file exists to patch the adalflow OllamaClient to support ollama batch embedding api `embed` + +from typing import Dict, Optional, Any + +from adalflow.components.model_client.ollama_client import ( + OllamaClient, + RequestError, + ResponseError, + log, +) +from adalflow.core.types import EmbedderOutput, Embedding + +from adalflow.core import ModelType + +import backoff + + +def convert_inputs_to_api_kwargs( + self, + input: Optional[Any] = None, + model_kwargs: Dict = {}, + model_type: ModelType = ModelType.UNDEFINED, +) -> Dict: + self.generate = False + final_model_kwargs = model_kwargs.copy() + if model_type == ModelType.EMBEDDER: + final_model_kwargs["input"] = input + return final_model_kwargs + elif model_type == ModelType.LLM: + if input is not None and input != "": + # check if "generate" is in model_kwargs, and if it is set to True, then we use generate api + if model_kwargs.get("generate", False): + final_model_kwargs["prompt"] = input + self.generate = True + else: + # for chat api, we need to convert the input to a message format + # if the input is a string, we create a message with role "user" + if isinstance(input, str): + input = [{"role": "user", "content": input}] + final_model_kwargs["messages"] = input + elif not isinstance(input, list): + raise ValueError("Input must be a string or a list of messages") + # if the input is a list of messages, we use it as is + return final_model_kwargs + else: + raise ValueError("Input must be text") + else: + raise ValueError(f"model_type {model_type} is not supported") + + +@backoff.on_exception( + backoff.expo, + (RequestError, ResponseError), + max_tries=5, +) +async def acall( + self, + api_kwargs: dict = None, + model_type: ModelType = ModelType.UNDEFINED, +): + if self.async_client is None: + self.init_async_client() + if self.async_client is None: + raise RuntimeError( + "Async client is not initialized" + ) + api_kwargs = api_kwargs or {} + if model_type == ModelType.EMBEDDER: + return await self.async_client.embed(**api_kwargs) + if model_type == ModelType.LLM: # in default we use chat + # create a message from the input + if "generate" in api_kwargs and api_kwargs["generate"]: + # remove generate from api_kwargs + api_kwargs.pop("generate") + return await self.async_client.generate(**api_kwargs) + else: + return await self.async_client.chat(**api_kwargs) + else: + raise ValueError(f"model_type {model_type} is not supported") + + + +@backoff.on_exception( + backoff.expo, + (RequestError, ResponseError), + max_tries=5, +) +def call( + self, + api_kwargs: dict = None, + model_type: ModelType = ModelType.UNDEFINED, +): + api_kwargs = api_kwargs or {} + if not self.sync_client: + self.init_sync_client() + if self.sync_client is None: + raise RuntimeError("Sync client is not initialized") + + if model_type == ModelType.EMBEDDER: + return self.sync_client.embed(**api_kwargs) + if model_type == ModelType.LLM: + if "generate" in api_kwargs and api_kwargs["generate"]: + # remove generate from api_kwargs + api_kwargs.pop("generate") + return self.sync_client.generate(**api_kwargs) + else: + return self.sync_client.chat(**api_kwargs) + else: + raise ValueError(f"model_type {model_type} is not supported") + + +def parse_embedding_response( + self, response: Dict[str, list[float]] +) -> EmbedderOutput: + r"""Parse the embedding response to a structure AdalFlow components can understand. + Pull the embedding from response['embedding'] and store it Embedding dataclass + """ + try: + return EmbedderOutput(data=[ + Embedding(embedding=emb, index=i) + for i, emb in enumerate(response["embeddings"]) + ]) + except Exception as e: + log.error(f"Error parsing the embedding response: {e}") + return EmbedderOutput(data=[], error=str(e), raw_response=response) + +OllamaClient.convert_inputs_to_api_kwargs = convert_inputs_to_api_kwargs +OllamaClient.parse_embedding_response = parse_embedding_response +OllamaClient.call = call +OllamaClient.acall = acall + diff --git a/api/rag.py b/api/rag.py index b85962b4f..3c0352fbe 100644 --- a/api/rag.py +++ b/api/rag.py @@ -189,22 +189,6 @@ def __init__(self, provider="google", model=None, use_s3: bool = False): # noqa # Initialize components self.memory = Memory() self.embedder = get_embedder(embedder_type=self.embedder_type) - - self_weakref = weakref.ref(self) - # Patch: ensure query embedding is always single string for Ollama - def single_string_embedder(query): - # Accepts either a string or a list, always returns embedding for a single string - if isinstance(query, list): - if len(query) != 1: - raise ValueError("Ollama embedder only supports a single string") - query = query[0] - instance = self_weakref() - assert instance is not None, "RAG instance is no longer available, but the query embedder was called." - return instance.embedder(input=query) - - # Use single string embedder for Ollama, regular embedder for others - self.query_embedder = single_string_embedder if self.is_ollama_embedder else self.embedder - self.initialize_db_manager() # Set up the output parser @@ -381,10 +365,9 @@ def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_ try: # Use the appropriate embedder for retrieval - retrieve_embedder = self.query_embedder if self.is_ollama_embedder else self.embedder self.retriever = FAISSRetriever( **configs["retriever"], - embedder=retrieve_embedder, + embedder=self.embedder, documents=self.transformed_docs, document_map_func=lambda doc: doc.vector, ) From 98e499e0bab65f49ac5d3cc362e7384818b8f519 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 18 Jul 2026 21:07:52 +0800 Subject: [PATCH 04/25] Revert "patch: patch adalflow OllamaClient to support batch embedding api. (#3)" (#4) This reverts commit 6bdc4e8ddb53a0ca705550b5eda17395bab34d83. --- api/config.py | 3 +- api/data_pipeline.py | 14 +++-- api/ollama_client.py | 131 ------------------------------------------- api/rag.py | 19 ++++++- 4 files changed, 29 insertions(+), 138 deletions(-) delete mode 100644 api/ollama_client.py diff --git a/api/config.py b/api/config.py index 4cafcdbc7..1d6a5874c 100644 --- a/api/config.py +++ b/api/config.py @@ -14,8 +14,7 @@ from api.google_embedder_client import GoogleEmbedderClient from api.azureai_client import AzureAIClient from api.dashscope_client import DashscopeClient -from api.ollama_client import OllamaClient -from adalflow import GoogleGenAIClient +from adalflow import GoogleGenAIClient, OllamaClient # Get API keys from environment variables OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') diff --git a/api/data_pipeline.py b/api/data_pipeline.py index 1d88ef9f7..f98068651 100644 --- a/api/data_pipeline.py +++ b/api/data_pipeline.py @@ -415,10 +415,16 @@ def prepare_data_pipeline(embedder_type: str = None, is_ollama_embedder: bool = embedder = get_embedder(embedder_type=embedder_type) - batch_size = embedder_config.get("batch_size", 500) - embedder_transformer = ToEmbeddings( - embedder=embedder, batch_size=batch_size - ) + # Choose appropriate processor based on embedder type + if embedder_type == 'ollama': + # Use Ollama document processor for single-document processing + embedder_transformer = OllamaDocumentProcessor(embedder=embedder) + else: + # Use batch processing for OpenAI and Google embedders + batch_size = embedder_config.get("batch_size", 500) + embedder_transformer = ToEmbeddings( + embedder=embedder, batch_size=batch_size + ) data_transformer = adal.Sequential( splitter, embedder_transformer diff --git a/api/ollama_client.py b/api/ollama_client.py deleted file mode 100644 index 2418e313d..000000000 --- a/api/ollama_client.py +++ /dev/null @@ -1,131 +0,0 @@ -# This file exists to patch the adalflow OllamaClient to support ollama batch embedding api `embed` - -from typing import Dict, Optional, Any - -from adalflow.components.model_client.ollama_client import ( - OllamaClient, - RequestError, - ResponseError, - log, -) -from adalflow.core.types import EmbedderOutput, Embedding - -from adalflow.core import ModelType - -import backoff - - -def convert_inputs_to_api_kwargs( - self, - input: Optional[Any] = None, - model_kwargs: Dict = {}, - model_type: ModelType = ModelType.UNDEFINED, -) -> Dict: - self.generate = False - final_model_kwargs = model_kwargs.copy() - if model_type == ModelType.EMBEDDER: - final_model_kwargs["input"] = input - return final_model_kwargs - elif model_type == ModelType.LLM: - if input is not None and input != "": - # check if "generate" is in model_kwargs, and if it is set to True, then we use generate api - if model_kwargs.get("generate", False): - final_model_kwargs["prompt"] = input - self.generate = True - else: - # for chat api, we need to convert the input to a message format - # if the input is a string, we create a message with role "user" - if isinstance(input, str): - input = [{"role": "user", "content": input}] - final_model_kwargs["messages"] = input - elif not isinstance(input, list): - raise ValueError("Input must be a string or a list of messages") - # if the input is a list of messages, we use it as is - return final_model_kwargs - else: - raise ValueError("Input must be text") - else: - raise ValueError(f"model_type {model_type} is not supported") - - -@backoff.on_exception( - backoff.expo, - (RequestError, ResponseError), - max_tries=5, -) -async def acall( - self, - api_kwargs: dict = None, - model_type: ModelType = ModelType.UNDEFINED, -): - if self.async_client is None: - self.init_async_client() - if self.async_client is None: - raise RuntimeError( - "Async client is not initialized" - ) - api_kwargs = api_kwargs or {} - if model_type == ModelType.EMBEDDER: - return await self.async_client.embed(**api_kwargs) - if model_type == ModelType.LLM: # in default we use chat - # create a message from the input - if "generate" in api_kwargs and api_kwargs["generate"]: - # remove generate from api_kwargs - api_kwargs.pop("generate") - return await self.async_client.generate(**api_kwargs) - else: - return await self.async_client.chat(**api_kwargs) - else: - raise ValueError(f"model_type {model_type} is not supported") - - - -@backoff.on_exception( - backoff.expo, - (RequestError, ResponseError), - max_tries=5, -) -def call( - self, - api_kwargs: dict = None, - model_type: ModelType = ModelType.UNDEFINED, -): - api_kwargs = api_kwargs or {} - if not self.sync_client: - self.init_sync_client() - if self.sync_client is None: - raise RuntimeError("Sync client is not initialized") - - if model_type == ModelType.EMBEDDER: - return self.sync_client.embed(**api_kwargs) - if model_type == ModelType.LLM: - if "generate" in api_kwargs and api_kwargs["generate"]: - # remove generate from api_kwargs - api_kwargs.pop("generate") - return self.sync_client.generate(**api_kwargs) - else: - return self.sync_client.chat(**api_kwargs) - else: - raise ValueError(f"model_type {model_type} is not supported") - - -def parse_embedding_response( - self, response: Dict[str, list[float]] -) -> EmbedderOutput: - r"""Parse the embedding response to a structure AdalFlow components can understand. - Pull the embedding from response['embedding'] and store it Embedding dataclass - """ - try: - return EmbedderOutput(data=[ - Embedding(embedding=emb, index=i) - for i, emb in enumerate(response["embeddings"]) - ]) - except Exception as e: - log.error(f"Error parsing the embedding response: {e}") - return EmbedderOutput(data=[], error=str(e), raw_response=response) - -OllamaClient.convert_inputs_to_api_kwargs = convert_inputs_to_api_kwargs -OllamaClient.parse_embedding_response = parse_embedding_response -OllamaClient.call = call -OllamaClient.acall = acall - diff --git a/api/rag.py b/api/rag.py index 3c0352fbe..b85962b4f 100644 --- a/api/rag.py +++ b/api/rag.py @@ -189,6 +189,22 @@ def __init__(self, provider="google", model=None, use_s3: bool = False): # noqa # Initialize components self.memory = Memory() self.embedder = get_embedder(embedder_type=self.embedder_type) + + self_weakref = weakref.ref(self) + # Patch: ensure query embedding is always single string for Ollama + def single_string_embedder(query): + # Accepts either a string or a list, always returns embedding for a single string + if isinstance(query, list): + if len(query) != 1: + raise ValueError("Ollama embedder only supports a single string") + query = query[0] + instance = self_weakref() + assert instance is not None, "RAG instance is no longer available, but the query embedder was called." + return instance.embedder(input=query) + + # Use single string embedder for Ollama, regular embedder for others + self.query_embedder = single_string_embedder if self.is_ollama_embedder else self.embedder + self.initialize_db_manager() # Set up the output parser @@ -365,9 +381,10 @@ def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_ try: # Use the appropriate embedder for retrieval + retrieve_embedder = self.query_embedder if self.is_ollama_embedder else self.embedder self.retriever = FAISSRetriever( **configs["retriever"], - embedder=self.embedder, + embedder=retrieve_embedder, documents=self.transformed_docs, document_map_func=lambda doc: doc.vector, ) From 7ef9f1985ef91ce42883f46e7de23deae9e7deeb Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 18 Jul 2026 21:13:22 +0800 Subject: [PATCH 05/25] Feat/ollama batch embed patch (#5) * patch: patch adalflow OllamaClient to support batch embedding api. related to https://github.com/SylphAI-Inc/AdalFlow/pull/496 * remove `OllamaDocumentProcessor` --- api/config.py | 3 +- api/data_pipeline.py | 15 ++--- api/ollama_client.py | 131 +++++++++++++++++++++++++++++++++++++++++++ api/ollama_patch.py | 51 ----------------- api/rag.py | 19 +------ 5 files changed, 138 insertions(+), 81 deletions(-) create mode 100644 api/ollama_client.py diff --git a/api/config.py b/api/config.py index 1d6a5874c..4cafcdbc7 100644 --- a/api/config.py +++ b/api/config.py @@ -14,7 +14,8 @@ from api.google_embedder_client import GoogleEmbedderClient from api.azureai_client import AzureAIClient from api.dashscope_client import DashscopeClient -from adalflow import GoogleGenAIClient, OllamaClient +from api.ollama_client import OllamaClient +from adalflow import GoogleGenAIClient # Get API keys from environment variables OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') diff --git a/api/data_pipeline.py b/api/data_pipeline.py index f98068651..9958f9262 100644 --- a/api/data_pipeline.py +++ b/api/data_pipeline.py @@ -11,7 +11,6 @@ from adalflow.utils import get_adalflow_default_root_path from adalflow.core.db import LocalDB from api.config import configs, DEFAULT_EXCLUDED_DIRS, DEFAULT_EXCLUDED_FILES -from api.ollama_patch import OllamaDocumentProcessor from urllib.parse import urlparse, urlunparse, quote import requests from requests.exceptions import RequestException @@ -415,16 +414,10 @@ def prepare_data_pipeline(embedder_type: str = None, is_ollama_embedder: bool = embedder = get_embedder(embedder_type=embedder_type) - # Choose appropriate processor based on embedder type - if embedder_type == 'ollama': - # Use Ollama document processor for single-document processing - embedder_transformer = OllamaDocumentProcessor(embedder=embedder) - else: - # Use batch processing for OpenAI and Google embedders - batch_size = embedder_config.get("batch_size", 500) - embedder_transformer = ToEmbeddings( - embedder=embedder, batch_size=batch_size - ) + batch_size = embedder_config.get("batch_size", 500) + embedder_transformer = ToEmbeddings( + embedder=embedder, batch_size=batch_size + ) data_transformer = adal.Sequential( splitter, embedder_transformer diff --git a/api/ollama_client.py b/api/ollama_client.py new file mode 100644 index 000000000..2418e313d --- /dev/null +++ b/api/ollama_client.py @@ -0,0 +1,131 @@ +# This file exists to patch the adalflow OllamaClient to support ollama batch embedding api `embed` + +from typing import Dict, Optional, Any + +from adalflow.components.model_client.ollama_client import ( + OllamaClient, + RequestError, + ResponseError, + log, +) +from adalflow.core.types import EmbedderOutput, Embedding + +from adalflow.core import ModelType + +import backoff + + +def convert_inputs_to_api_kwargs( + self, + input: Optional[Any] = None, + model_kwargs: Dict = {}, + model_type: ModelType = ModelType.UNDEFINED, +) -> Dict: + self.generate = False + final_model_kwargs = model_kwargs.copy() + if model_type == ModelType.EMBEDDER: + final_model_kwargs["input"] = input + return final_model_kwargs + elif model_type == ModelType.LLM: + if input is not None and input != "": + # check if "generate" is in model_kwargs, and if it is set to True, then we use generate api + if model_kwargs.get("generate", False): + final_model_kwargs["prompt"] = input + self.generate = True + else: + # for chat api, we need to convert the input to a message format + # if the input is a string, we create a message with role "user" + if isinstance(input, str): + input = [{"role": "user", "content": input}] + final_model_kwargs["messages"] = input + elif not isinstance(input, list): + raise ValueError("Input must be a string or a list of messages") + # if the input is a list of messages, we use it as is + return final_model_kwargs + else: + raise ValueError("Input must be text") + else: + raise ValueError(f"model_type {model_type} is not supported") + + +@backoff.on_exception( + backoff.expo, + (RequestError, ResponseError), + max_tries=5, +) +async def acall( + self, + api_kwargs: dict = None, + model_type: ModelType = ModelType.UNDEFINED, +): + if self.async_client is None: + self.init_async_client() + if self.async_client is None: + raise RuntimeError( + "Async client is not initialized" + ) + api_kwargs = api_kwargs or {} + if model_type == ModelType.EMBEDDER: + return await self.async_client.embed(**api_kwargs) + if model_type == ModelType.LLM: # in default we use chat + # create a message from the input + if "generate" in api_kwargs and api_kwargs["generate"]: + # remove generate from api_kwargs + api_kwargs.pop("generate") + return await self.async_client.generate(**api_kwargs) + else: + return await self.async_client.chat(**api_kwargs) + else: + raise ValueError(f"model_type {model_type} is not supported") + + + +@backoff.on_exception( + backoff.expo, + (RequestError, ResponseError), + max_tries=5, +) +def call( + self, + api_kwargs: dict = None, + model_type: ModelType = ModelType.UNDEFINED, +): + api_kwargs = api_kwargs or {} + if not self.sync_client: + self.init_sync_client() + if self.sync_client is None: + raise RuntimeError("Sync client is not initialized") + + if model_type == ModelType.EMBEDDER: + return self.sync_client.embed(**api_kwargs) + if model_type == ModelType.LLM: + if "generate" in api_kwargs and api_kwargs["generate"]: + # remove generate from api_kwargs + api_kwargs.pop("generate") + return self.sync_client.generate(**api_kwargs) + else: + return self.sync_client.chat(**api_kwargs) + else: + raise ValueError(f"model_type {model_type} is not supported") + + +def parse_embedding_response( + self, response: Dict[str, list[float]] +) -> EmbedderOutput: + r"""Parse the embedding response to a structure AdalFlow components can understand. + Pull the embedding from response['embedding'] and store it Embedding dataclass + """ + try: + return EmbedderOutput(data=[ + Embedding(embedding=emb, index=i) + for i, emb in enumerate(response["embeddings"]) + ]) + except Exception as e: + log.error(f"Error parsing the embedding response: {e}") + return EmbedderOutput(data=[], error=str(e), raw_response=response) + +OllamaClient.convert_inputs_to_api_kwargs = convert_inputs_to_api_kwargs +OllamaClient.parse_embedding_response = parse_embedding_response +OllamaClient.call = call +OllamaClient.acall = acall + diff --git a/api/ollama_patch.py b/api/ollama_patch.py index fb985dfaf..bb6a77294 100644 --- a/api/ollama_patch.py +++ b/api/ollama_patch.py @@ -1,10 +1,4 @@ -from typing import Sequence, List -from copy import deepcopy -from tqdm import tqdm import logging -import adalflow as adal -from adalflow.core.types import Document -from adalflow.core.component import DataComponent import requests import os @@ -58,48 +52,3 @@ def check_ollama_model_exists(model_name: str, ollama_host: str = None) -> bool: except Exception as e: logger.warning(f"Error checking Ollama model availability: {e}") return False - -class OllamaDocumentProcessor(DataComponent): - """ - Process documents for Ollama embeddings by processing one document at a time. - Adalflow Ollama Client does not support batch embedding, so we need to process each document individually. - """ - def __init__(self, embedder: adal.Embedder) -> None: - super().__init__() - self.embedder = embedder - - def __call__(self, documents: Sequence[Document]) -> Sequence[Document]: - output = deepcopy(documents) - logger.info(f"Processing {len(output)} documents individually for Ollama embeddings") - - successful_docs = [] - expected_embedding_size = None - - for i, doc in enumerate(tqdm(output, desc="Processing documents for Ollama embeddings")): - try: - # Get embedding for a single document - result = self.embedder(input=doc.text) - if result.data and len(result.data) > 0: - embedding = result.data[0].embedding - - # Validate embedding size consistency - if expected_embedding_size is None: - expected_embedding_size = len(embedding) - logger.info(f"Expected embedding size set to: {expected_embedding_size}") - elif len(embedding) != expected_embedding_size: - file_path = getattr(doc, 'meta_data', {}).get('file_path', f'document_{i}') - logger.warning(f"Document '{file_path}' has inconsistent embedding size {len(embedding)} != {expected_embedding_size}, skipping") - continue - - # Assign the embedding to the document - output[i].vector = embedding - successful_docs.append(output[i]) - else: - file_path = getattr(doc, 'meta_data', {}).get('file_path', f'document_{i}') - logger.warning(f"Failed to get embedding for document '{file_path}', skipping") - except Exception as e: - file_path = getattr(doc, 'meta_data', {}).get('file_path', f'document_{i}') - logger.error(f"Error processing document '{file_path}': {e}, skipping") - - logger.info(f"Successfully processed {len(successful_docs)}/{len(output)} documents with consistent embeddings") - return successful_docs \ No newline at end of file diff --git a/api/rag.py b/api/rag.py index b85962b4f..3c0352fbe 100644 --- a/api/rag.py +++ b/api/rag.py @@ -189,22 +189,6 @@ def __init__(self, provider="google", model=None, use_s3: bool = False): # noqa # Initialize components self.memory = Memory() self.embedder = get_embedder(embedder_type=self.embedder_type) - - self_weakref = weakref.ref(self) - # Patch: ensure query embedding is always single string for Ollama - def single_string_embedder(query): - # Accepts either a string or a list, always returns embedding for a single string - if isinstance(query, list): - if len(query) != 1: - raise ValueError("Ollama embedder only supports a single string") - query = query[0] - instance = self_weakref() - assert instance is not None, "RAG instance is no longer available, but the query embedder was called." - return instance.embedder(input=query) - - # Use single string embedder for Ollama, regular embedder for others - self.query_embedder = single_string_embedder if self.is_ollama_embedder else self.embedder - self.initialize_db_manager() # Set up the output parser @@ -381,10 +365,9 @@ def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_ try: # Use the appropriate embedder for retrieval - retrieve_embedder = self.query_embedder if self.is_ollama_embedder else self.embedder self.retriever = FAISSRetriever( **configs["retriever"], - embedder=retrieve_embedder, + embedder=self.embedder, documents=self.transformed_docs, document_map_func=lambda doc: doc.vector, ) From 6d2869e9e676dba3a3f153db1f045dbacbca60b0 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 18 Jul 2026 21:18:52 +0800 Subject: [PATCH 06/25] Refactor/chat streamer (#6) * refactor chat streaming * update `_is_token_limit_error` logic * test: add chat provider unittests * move `_is_token_limit_error` to module level * remove unnecessary TYPE_CHECKING * add EOF * update prompt function * refactor `websocket_wiki.py` using `ChatStreamer` * remove unnecessary imports * add litellm test * update kwargs in `GoogleGenerativeChatStreamer` --- api/chat/__init__.py | 20 ++ api/chat/_prompts.py | 34 +++ api/chat/_stream.py | 315 +++++++++++++++++++++ api/simple_chat.py | 80 ++---- api/websocket_wiki.py | 604 ++++------------------------------------ tests/unit/test_chat.py | 7 +- 6 files changed, 455 insertions(+), 605 deletions(-) create mode 100644 api/chat/__init__.py create mode 100644 api/chat/_prompts.py create mode 100644 api/chat/_stream.py diff --git a/api/chat/__init__.py b/api/chat/__init__.py new file mode 100644 index 000000000..57008ffc1 --- /dev/null +++ b/api/chat/__init__.py @@ -0,0 +1,20 @@ +from api.chat._stream import ChatStreamer +from api.chat._prompts import prompt_builder + + +def is_token_limit_error(exc: Exception) -> bool: + error_message = str(exc).lower() + return any( + k in error_message for k in ( + "maximum context length", + "token limit", + "too many tokens", + ) + ) + + +__all__ = [ + "ChatStreamer", + "prompt_builder", + "is_token_limit_error", +] diff --git a/api/chat/_prompts.py b/api/chat/_prompts.py new file mode 100644 index 000000000..bef5254d6 --- /dev/null +++ b/api/chat/_prompts.py @@ -0,0 +1,34 @@ +from api.logging_config import setup_logging +from logging import getLogger + +setup_logging() +logger = getLogger(__name__) + + +def prompt_builder( + system_prompt: str, + query: str, + conversation_history: str | None = None, + file_path: str | None = None, + file_content: str | None = None, + context: str = "", + simplify: bool = False, +) -> str: + prompt = f"/no_think {system_prompt}\n\n" + if conversation_history: + prompt += f"\n{conversation_history}\n\n" + + if file_path and file_content: + prompt += f"\n{file_content}\n\n\n" + + if not simplify: + if context.strip(): + context_prompt = f"\n{context}\n\n\n" + else: + # Add a note that we're skipping RAG due to size constraints or because it's the isolated API + logger.info("No context available from RAG") + context_prompt = "Answering without retrieval augmentation.\n\n" + else: + context_prompt = "Answering without retrieval augmentation due to input size constraints.\n\n" + + return prompt + context_prompt + f"\n{query}\n\n\nAssistant: " diff --git a/api/chat/_stream.py b/api/chat/_stream.py new file mode 100644 index 000000000..94251d162 --- /dev/null +++ b/api/chat/_stream.py @@ -0,0 +1,315 @@ +import logging +from abc import abstractmethod, ABC + +from typing import TYPE_CHECKING + +from collections.abc import AsyncIterator + +from adalflow.core.types import ModelType +from api.config import ( + OPENROUTER_API_KEY, + OPENAI_API_KEY, + AWS_ACCESS_KEY_ID, + AWS_SECRET_ACCESS_KEY, + LITELLM_API_KEY, +) + +if TYPE_CHECKING: + from ollama import ChatResponse + from openai.types.chat import ChatCompletionChunk + from openai import AsyncStream + from api.openai_client import OpenAIClient + +MODEL_CFG = dict[str, str | int | float] + +# Configure logging +from api.logging_config import setup_logging + +setup_logging() +logger = logging.getLogger(__name__) + + +class ChatStreamer(ABC): + _registry: dict[str, type["ChatStreamer"]] = {} + provider: str + error_hint: str | None = None + + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + if provider := getattr(cls, "provider", None): + ChatStreamer._registry[provider] = cls + + @classmethod + def create(cls, *, provider: str, model: str | None = None, model_config: MODEL_CFG) -> "ChatStreamer": + model = model or model_config.get("model") + logger.info("Using %s with model: %s", provider, model) + registered = ChatStreamer._registry.get(provider, None) + if registered: + return registered(model=model, model_config=model_config) + raise RuntimeError(f"Provider {provider} not registered") + + @abstractmethod + def respond_stream(self, prompt: str) -> AsyncIterator[str]: + raise NotImplementedError(f"{type(self).__name__} does not implement `respond_stream`") + + +class OllamaChatStreamer(ChatStreamer): + provider = "ollama" + + def __init__(self, *, model: str, model_config: MODEL_CFG): + from adalflow.components.model_client.ollama_client import OllamaClient + + self.client = OllamaClient() + self.model_kwargs = { + "model": model, + "stream": True, + "options": { + "temperature": model_config["temperature"], + "top_p": model_config["top_p"], + "num_ctx": model_config["num_ctx"] + } + } + + logger.debug(f"Prompting Ollama with kwargs: {self.model_kwargs}") + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt + " /no_think", # todo I think this could be added into model kwargs? + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM, + ) + + response: "AsyncIterator[ChatResponse]" = await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ) + + async for chunk in response: + if not hasattr(chunk, "message"): + raise RuntimeError( + "`message` field not found in response. Wrong ollama-python version probably.", + ) + text = chunk.message.content + if text: + text = text.replace('', '').replace('', '') + yield text + + +class OpenRouterChatStreamer(ChatStreamer): + provider = "openrouter" + error_hint = ( + "Please check that you have set the OPENROUTER_API_KEY " + "environment variable with a valid API key." + ) + + def __init__(self, *, model: str, model_config: MODEL_CFG): + if not OPENROUTER_API_KEY: + logger.warning("OPENROUTER_API_KEY not configured, but continuing with request") + # We'll let the OpenRouterClient handle this and return a friendly error message + from api.openrouter_client import OpenRouterClient + + self.client = OpenRouterClient() + self.model_kwargs = { + "model": model, + "stream": True, + "temperature": model_config["temperature"] + } + if "top_k" in model_config: + self.model_kwargs["top_k"] = model_config["top_k"] + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt, + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM, + ) + async for chunk in await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ): + yield chunk + + +class _OpenAICompatStreamer(ChatStreamer): + client: "OpenAIClient" + model_kwargs: dict + + def __init__(self, *, model: str, model_config: MODEL_CFG): + self.client = self._build_client() + self.model_kwargs = { + "model": model, + "stream": True, + "temperature": model_config["temperature"] + } + # Only add top_p if it exists in the model config + if "top_p" in model_config: + self.model_kwargs["top_p"] = model_config["top_p"] + + @abstractmethod + def _build_client(self) -> "OpenAIClient": + raise NotImplementedError( + f"{type(self).__name__} must return an `OpenAIClient` instance" + ) + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt, + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM + ) + response: "AsyncStream[ChatCompletionChunk]" = await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ) + + async for chunk in response: + if ( + chunk.choices and + chunk.choices[0].delta is not None and + chunk.choices[0].delta.content is not None + ): + yield chunk.choices[0].delta.content + + +class OpenAIChatStreamer(_OpenAICompatStreamer): + provider = "openai" + error_hint = ( + "Please check that you have set the OPENAI_API_KEY " + "environment variable with a valid API key." + ) + + def __init__(self, *, model: str, model_config: MODEL_CFG): + if not OPENAI_API_KEY: + logger.warning("OPENAI_API_KEY not configured, but continuing with request") + + super().__init__(model=model, model_config=model_config) + + def _build_client(self): + from api.openai_client import OpenAIClient + return OpenAIClient() + + +class AzureChatStreamer(_OpenAICompatStreamer): + provider = "azure" + error_hint = ( + "Please check that you have set the AZURE_OPENAI_API_KEY, " + "AZURE_OPENAI_ENDPOINT, and AZURE_OPENAI_VERSION " + "environment variables with valid values." + ) + + def _build_client(self): + from api.azureai_client import AzureAIClient + return AzureAIClient() + + +class LiteLLMChatStreamer(_OpenAICompatStreamer): + provider = "litellm" + error_hint = ( + "Please check that you have set the LITELLM_API_KEY " + "environment variable with a valid API key." + ) + + def __init__(self, *, model: str, model_config: MODEL_CFG): + if not LITELLM_API_KEY: + logger.warning("LITELLM_API_KEY not configured, but continuing with request") + # We'll let the OpenAIClient handle this and return an error message + + super().__init__(model=model, model_config=model_config) + + def _build_client(self): + from api.litellm_client import LiteLLMClient + return LiteLLMClient() + + +class BedrockChatStreamer(ChatStreamer): + provider = "bedrock" + error_hint = ( + "Please check that you have set the AWS_ACCESS_KEY_ID " + "and AWS_SECRET_ACCESS_KEY environment variables with valid credentials." + ) + + def __init__(self, *, model: str, model_config: MODEL_CFG): + if not AWS_ACCESS_KEY_ID or not AWS_SECRET_ACCESS_KEY: + logger.warning("AWS_ACCESS_KEY_ID or AWS_SECRET_ACCESS_KEY not configured, but continuing with request") + # We'll let the BedrockClient handle this and return an error message + from api.bedrock_client import BedrockClient + + self.client = BedrockClient() + self.model_kwargs = {"model": model} + + for key in ( + "temperature", + "top_p", + ): + if key in model_config: + self.model_kwargs[key] = model_config[key] + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt, + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM, + ) + response = await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ) + if not isinstance(response, str): + response = str(response) + yield response + + +class DashScopeChatStreamer(ChatStreamer): + provider = "dashscope" + error_hint = ( + "Please check that you have set the DASHSCOPE_API_KEY (and optionally " + "DASHSCOPE_WORKSPACE_ID) environment variables with valid values." + ) + + def __init__(self, *, model: str, model_config: MODEL_CFG): + from api.dashscope_client import DashscopeClient + + self.client = DashscopeClient() + self.model_kwargs = { + "model": model, + "stream": True, + "temperature": model_config["temperature"], + "top_p": model_config["top_p"], + } + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + api_kwargs = self.client.convert_inputs_to_api_kwargs( + input=prompt, + model_kwargs=self.model_kwargs, + model_type=ModelType.LLM, + ) + response = await self.client.acall( + api_kwargs=api_kwargs, + model_type=ModelType.LLM, + ) + async for text in response: + if text: + yield text + + +class GoogleGenerativeChatStreamer(ChatStreamer): + provider = "google" + + def __init__(self, *, model: str, model_config: MODEL_CFG): + import google.generativeai as genai + from google.generativeai.types import GenerationConfig + + self.client = genai.GenerativeModel( + model_name=model, + generation_config=GenerationConfig( + temperature=model_config.get("temperature"), + top_p=model_config.get("top_p"), + top_k=model_config.get("top_k"), + ) + ) + + async def respond_stream(self, prompt: str) -> AsyncIterator[str]: + response = await self.client.generate_content_async(prompt, stream=True) + async for chunk in response: + if hasattr(chunk, "text"): + yield chunk.text diff --git a/api/simple_chat.py b/api/simple_chat.py index fb7fe4d07..ab1d9de39 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -1,13 +1,14 @@ import logging from typing import List, Optional, Callable from urllib.parse import unquote +from functools import partial from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field -from api.chat import ChatStreamer +from api.chat import ChatStreamer, , prompt_builder, is_token_limit_error from api.config import get_model_config, configs from api.data_pipeline import count_tokens, get_file_content from api.rag import RAG, MAX_INPUT_TOKENS @@ -66,13 +67,6 @@ class ChatCompletionRequest(BaseModel): included_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to include exclusively") -def _is_token_limit_error(exc: Exception) -> bool: - error_message = str(exc).lower() - return any( - k in error_message for k in ("maximum context length", "token limit", "too many tokens") - ) - - @app.post("/chat/completions/stream") async def chat_completions_stream(request: ChatCompletionRequest): """Stream a chat completion response directly using Google Generative AI""" @@ -304,43 +298,6 @@ async def chat_completions_stream(request: ChatCompletionRequest): if not isinstance(turn_id, int) and hasattr(turn, 'user_query') and hasattr(turn, 'assistant_response'): conversation_history += f"\n{turn.user_query.query_str}\n{turn.assistant_response.response_str}\n\n" - def prompt() -> str: - # Create the prompt with context - prompt = f"/no_think {system_prompt}\n\n" - if conversation_history: - prompt += f"\n{conversation_history}\n\n" - - # Check if filePath is provided and fetch file content if it exists - if request.filePath and file_content: - # Add file content to the prompt after conversation history - prompt += f"\n{file_content}\n\n\n" - - # Only include context if it's not empty - if context_text.strip(): - context_prompt = f"\n{context_text}\n\n\n" - else: - # Add a note that we're skipping RAG due to size constraints or because it's the isolated API - logger.info("No context available from RAG") - context_prompt = "Answering without retrieval augmentation.\n\n" - - prompt += f"{context_prompt}\n{query}\n\n\nAssistant: " - return prompt - - def simplified_prompt() -> str: - prompt = f"/no_think {system_prompt}\n\n" - if conversation_history: - prompt += f"\n{conversation_history}\n\n" - - # Include file content in the fallback prompt if it was retrieved - if request.filePath and file_content: - prompt += f"\n{file_content}\n\n\n" - - prompt += ( - f"Answering without retrieval augmentation due to input size constraints.\n\n" - f"\n{query}\n\n\nAssistant: " - ) - return prompt - async def stream_and_fallback( streamer: ChatStreamer, prompt_func: Callable[[], str], @@ -350,19 +307,23 @@ async def stream_and_fallback( async for chunk in streamer.respond_stream(prompt_func()): yield chunk except Exception as e: - if _is_token_limit_error(e): + if is_token_limit_error(e): logger.warning("Token limit exceeded, retrying without context") try: async for chunk in streamer.respond_stream(simplified_prompt_func()): yield chunk except Exception as e2: - logger.error(f"Error in fallback streaming response: {str(e2)}") + logger.error("Error in fallback streaming response: %s", str(e2)) yield ( f"\nI apologize, but your request is too large for me to process. " f"Please try a shorter query or break it into smaller parts." ) else: - yield f"\nError: {str(e)}" + error_str = f"Error with {streamer.provider} API: {e}" + logger.error(error_str, exc_info=True) + if streamer.error_hint: + error_str += f"\n\n{streamer.error_hint}" + yield "\n" + error_str model_config = get_model_config(request.provider, request.model)["model_kwargs"] chat_streamer = ChatStreamer.create( @@ -370,12 +331,31 @@ async def stream_and_fallback( model=request.model, model_config=model_config, ) + prompt_kwargs = dict( + system_prompt=system_prompt, + query=query, + conversation_history=conversation_history, + file_path=request.filePath, + file_content=file_content, + context=context_text, + ) + + prompt_func = partial( + prompt_builder, + **prompt_kwargs, + simplify=False, + ) + simplified_prompt_func = partial( + prompt_builder, + **prompt_kwargs, + simplify=True, + ) # Return streaming response return StreamingResponse(stream_and_fallback( streamer=chat_streamer, - prompt_func=prompt, - simplified_prompt_func=simplified_prompt, + prompt_func=prompt_func, + simplified_prompt_func=simplified_prompt_func, ), media_type="text/event-stream") except HTTPException: diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index 7064f54b4..87d508726 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -1,30 +1,18 @@ import logging -import os -from typing import List, Optional, Dict, Any +from collections.abc import AsyncIterator, Callable +from typing import List, Optional from urllib.parse import unquote +from functools import partial -import google.generativeai as genai -from adalflow.components.model_client.ollama_client import OllamaClient -from adalflow.core.types import ModelType -from fastapi import WebSocket, WebSocketDisconnect, HTTPException +from fastapi import WebSocket, WebSocketDisconnect from pydantic import BaseModel, Field +from api.chat import ChatStreamer, prompt_builder, is_token_limit_error from api.config import ( get_model_config, configs, - OPENROUTER_API_KEY, - OPENAI_API_KEY, - LITELLM_API_KEY, - AWS_ACCESS_KEY_ID, - AWS_SECRET_ACCESS_KEY, ) from api.data_pipeline import count_tokens, get_file_content -from api.bedrock_client import BedrockClient -from api.openai_client import OpenAIClient -from api.litellm_client import LiteLLMClient -from api.openrouter_client import OpenRouterClient -from api.azureai_client import AzureAIClient -from api.dashscope_client import DashscopeClient from api.rag import RAG # Configure logging @@ -439,544 +427,54 @@ async def handle_websocket_chat(websocket: WebSocket): prompt += f"\n{query}\n\n\nAssistant: " - model_config = get_model_config(request.provider, request.model)["model_kwargs"] - - if request.provider == "ollama": - prompt += " /no_think" - - model = OllamaClient() - model_kwargs = { - "model": model_config["model"], - "stream": True, - "options": { - "temperature": model_config["temperature"], - "top_p": model_config["top_p"], - "num_ctx": model_config["num_ctx"] - } - } - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "openrouter": - logger.info(f"Using OpenRouter with model: {request.model}") - - # Check if OpenRouter API key is set - if not OPENROUTER_API_KEY: - logger.warning("OPENROUTER_API_KEY not configured, but continuing with request") - # We'll let the OpenRouterClient handle this and return a friendly error message - - model = OpenRouterClient() - model_kwargs = { - "model": request.model, - "stream": True, - "temperature": model_config["temperature"] - } - # Only add top_p if it exists in the model config - if "top_p" in model_config: - model_kwargs["top_p"] = model_config["top_p"] - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "openai": - logger.info(f"Using Openai protocol with model: {request.model}") - - # Check if an API key is set for Openai - if not OPENAI_API_KEY: - logger.warning("OPENAI_API_KEY not configured, but continuing with request") - # We'll let the OpenAIClient handle this and return an error message - - # Initialize Openai client - model = OpenAIClient() - model_kwargs = { - "model": request.model, - "stream": True, - "temperature": model_config["temperature"] - } - # Only add top_p if it exists in the model config - if "top_p" in model_config: - model_kwargs["top_p"] = model_config["top_p"] - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "litellm": - logger.info(f"Using Openai protocol with model on LiteLLM: {request.model}") - - # Check if an API key is set for Litellm - if not LITELLM_API_KEY: - logger.warning("LITELLM_API_KEY not configured, but continuing with request") - # We'll let the OpenAIClient handle this and return an error message - - # Initialize LiteLLM client - model = LiteLLMClient() - model_kwargs = { - "model": request.model, - "stream": True, - "temperature": model_config["temperature"] - } - # Only add top_p if it exists in the model config - if "top_p" in model_config: - model_kwargs["top_p"] = model_config["top_p"] - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "bedrock": - logger.info(f"Using AWS Bedrock with model: {request.model}") - - if not AWS_ACCESS_KEY_ID or not AWS_SECRET_ACCESS_KEY: - logger.warning( - "AWS_ACCESS_KEY_ID or AWS_SECRET_ACCESS_KEY not configured, but continuing with request") - - model = BedrockClient() - model_kwargs = { - "model": request.model, - } - - for key in ["temperature", "top_p"]: - if key in model_config: - model_kwargs[key] = model_config[key] - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "azure": - logger.info(f"Using Azure AI with model: {request.model}") - - # Initialize Azure AI client - model = AzureAIClient() - model_kwargs = { - "model": request.model, - "stream": True, - "temperature": model_config["temperature"], - "top_p": model_config["top_p"] - } - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - elif request.provider == "dashscope": - logger.info(f"Using Dashscope with model: {request.model}") - - # Initialize Dashscope client - model = DashscopeClient() - model_kwargs = { - "model": request.model, - "stream": True, - "temperature": model_config["temperature"], - "top_p": model_config["top_p"] - } - - api_kwargs = model.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - else: - # Initialize Google Generative AI model - model = genai.GenerativeModel( - model_name=model_config["model"], - generation_config={ - "temperature": model_config["temperature"], - "top_p": model_config["top_p"], - "top_k": model_config["top_k"] - } - ) - - # Process the response based on the provider - try: - if request.provider == "ollama": - # Get the response and handle it properly using the previously created api_kwargs - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - # Handle streaming response from Ollama - async for chunk in response: - text = None - if isinstance(chunk, dict): - text = chunk.get("message", {}).get("content") if isinstance(chunk.get("message"), dict) else chunk.get("message") - else: - message = getattr(chunk, "message", None) - if message is not None: - if isinstance(message, dict): - text = message.get("content") - else: - text = getattr(message, "content", None) - - if not text: - text = getattr(chunk, 'response', None) or getattr(chunk, 'text', None) - - if not text and hasattr(chunk, "__dict__"): - message = chunk.__dict__.get("message") - if isinstance(message, dict): - text = message.get("content") - - if isinstance(text, str) and text and not text.startswith('model=') and not text.startswith('created_at='): - clean_text = text.replace('', '').replace('', '') - await websocket.send_text(clean_text) - # Explicitly close the WebSocket connection after the response is complete - await websocket.close() - elif request.provider == "openrouter": - try: - # Get the response and handle it properly using the previously created api_kwargs - logger.info("Making OpenRouter API call") - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - # Handle streaming response from OpenRouter - async for chunk in response: - await websocket.send_text(chunk) - # Explicitly close the WebSocket connection after the response is complete - await websocket.close() - except Exception as e_openrouter: - logger.error(f"Error with OpenRouter API: {str(e_openrouter)}") - error_msg = f"\nError with OpenRouter API: {str(e_openrouter)}\n\nPlease check that you have set the OPENROUTER_API_KEY environment variable with a valid API key." - await websocket.send_text(error_msg) - # Close the WebSocket connection after sending the error message - await websocket.close() - elif request.provider == "openai": - try: - # Get the response and handle it properly using the previously created api_kwargs - logger.info("Making Openai API call") - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - # Handle streaming response from Openai - async for chunk in response: - choices = getattr(chunk, "choices", []) - if len(choices) > 0: - delta = getattr(choices[0], "delta", None) - if delta is not None: - text = getattr(delta, "content", None) - if text is not None: - await websocket.send_text(text) - # Explicitly close the WebSocket connection after the response is complete - await websocket.close() - except Exception as e_openai: - logger.error(f"Error with Openai API: {str(e_openai)}") - error_msg = f"\nError with Openai API: {str(e_openai)}\n\nPlease check that you have set the OPENAI_API_KEY environment variable with a valid API key." - await websocket.send_text(error_msg) - # Close the WebSocket connection after sending the error message - await websocket.close() - elif request.provider == "litellm": - try: - # Get the response and handle it properly using the previously created api_kwargs - logger.info("Making LiteLLM API call") - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - # Handle streaming response from LiteLLM - async for chunk in response: - choices = getattr(chunk, "choices", []) - if len(choices) > 0: - delta = getattr(choices[0], "delta", None) - if delta is not None: - text = getattr(delta, "content", None) - if text is not None: - await websocket.send_text(text) - # Explicitly close the WebSocket connection after the response is complete - await websocket.close() - except Exception as e_litellm: - logger.error(f"Error with LiteLLM API: {str(e_litellm)}") - error_msg = f"\nError with LiteLLM API: {str(e_litellm)}\n\nPlease check that you have set the LITELLM_API_KEY environment variable with a valid API key." - await websocket.send_text(error_msg) - # Close the WebSocket connection after sending the error message - await websocket.close() - elif request.provider == "bedrock": - try: - logger.info("Making AWS Bedrock API call") - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - if isinstance(response, str): - await websocket.send_text(response) - else: - await websocket.send_text(str(response)) - await websocket.close() - except Exception as e_bedrock: - logger.error(f"Error with AWS Bedrock API: {str(e_bedrock)}") - error_msg = ( - f"\nError with AWS Bedrock API: {str(e_bedrock)}\n\n" - "Please check that you have set the AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY " - "environment variables with valid credentials." - ) - await websocket.send_text(error_msg) - await websocket.close() - elif request.provider == "azure": - try: - # Get the response and handle it properly using the previously created api_kwargs - logger.info("Making Azure AI API call") - response = await model.acall(api_kwargs=api_kwargs, model_type=ModelType.LLM) - # Handle streaming response from Azure AI - async for chunk in response: - choices = getattr(chunk, "choices", []) - if len(choices) > 0: - delta = getattr(choices[0], "delta", None) - if delta is not None: - text = getattr(delta, "content", None) - if text is not None: - await websocket.send_text(text) - # Explicitly close the WebSocket connection after the response is complete - await websocket.close() - except Exception as e_azure: - logger.error(f"Error with Azure AI API: {str(e_azure)}") - error_msg = f"\nError with Azure AI API: {str(e_azure)}\n\nPlease check that you have set the AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, and AZURE_OPENAI_VERSION environment variables with valid values." - await websocket.send_text(error_msg) - # Close the WebSocket connection after sending the error message - await websocket.close() - elif request.provider == "dashscope": - try: - # Get the response and handle it properly using the previously created api_kwargs - logger.info("Making Dashscope API call") - response = await model.acall( - api_kwargs=api_kwargs, model_type=ModelType.LLM - ) - # DashscopeClient.acall with stream=True returns an async - # generator of plain text chunks - async for text in response: - if text: - await websocket.send_text(text) - # Explicitly close the WebSocket connection after the response is complete - await websocket.close() - except Exception as e_dashscope: - logger.error(f"Error with Dashscope API: {str(e_dashscope)}") - error_msg = ( - f"\nError with Dashscope API: {str(e_dashscope)}\n\n" - "Please check that you have set the DASHSCOPE_API_KEY (and optionally " - "DASHSCOPE_WORKSPACE_ID) environment variables with valid values." - ) - await websocket.send_text(error_msg) - # Close the WebSocket connection after sending the error message - await websocket.close() - else: - # Google Generative AI (default provider) - response = model.generate_content(prompt, stream=True) - for chunk in response: - if hasattr(chunk, 'text'): - await websocket.send_text(chunk.text) - await websocket.close() - - except Exception as e_outer: - logger.error(f"Error in streaming response: {str(e_outer)}") - error_message = str(e_outer) - - # Check for token limit errors - if "maximum context length" in error_message or "token limit" in error_message or "too many tokens" in error_message: - # If we hit a token limit error, try again without context - logger.warning("Token limit exceeded, retrying without context") - try: - # Create a simplified prompt without context - simplified_prompt = f"/no_think {system_prompt}\n\n" - if conversation_history: - simplified_prompt += f"\n{conversation_history}\n\n" - - # Include file content in the fallback prompt if it was retrieved - if request.filePath and file_content: - simplified_prompt += f"\n{file_content}\n\n\n" - - simplified_prompt += "Answering without retrieval augmentation due to input size constraints.\n\n" - simplified_prompt += f"\n{query}\n\n\nAssistant: " - - if request.provider == "ollama": - simplified_prompt += " /no_think" - - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - fallback_response = await model.acall(api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM) - - # Handle streaming fallback_response from Ollama - async for chunk in fallback_response: - text = getattr(chunk, 'response', None) or getattr(chunk, 'text', None) or str(chunk) - if text and not text.startswith('model=') and not text.startswith('created_at='): - text = text.replace('', '').replace('', '') - await websocket.send_text(text) - elif request.provider == "openrouter": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - logger.info("Making fallback OpenRouter API call") - fallback_response = await model.acall(api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM) - - # Handle streaming fallback_response from OpenRouter - async for chunk in fallback_response: - await websocket.send_text(chunk) - except Exception as e_fallback: - logger.error(f"Error with OpenRouter API fallback: {str(e_fallback)}") - error_msg = f"\nError with OpenRouter API fallback: {str(e_fallback)}\n\nPlease check that you have set the OPENROUTER_API_KEY environment variable with a valid API key." - await websocket.send_text(error_msg) - elif request.provider == "openai": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - logger.info("Making fallback Openai API call") - fallback_response = await model.acall(api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM) - - # Handle streaming fallback_response from Openai - async for chunk in fallback_response: - text = chunk if isinstance(chunk, str) else getattr(chunk, 'text', str(chunk)) - await websocket.send_text(text) - except Exception as e_fallback: - logger.error(f"Error with Openai API fallback: {str(e_fallback)}") - error_msg = f"\nError with Openai API fallback: {str(e_fallback)}\n\nPlease check that you have set the OPENAI_API_KEY environment variable with a valid API key." - await websocket.send_text(error_msg) - elif request.provider == "litellm": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - logger.info("Making fallback LiteLLM API call") - fallback_response = await model.acall(api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM) - - # Handle streaming fallback_response from LiteLLM - async for chunk in fallback_response: - choices = getattr(chunk, "choices", []) - if len(choices) > 0: - delta = getattr(choices[0], "delta", None) - if delta is not None: - text = getattr(delta, "content", None) - if text is not None: - await websocket.send_text(text) - except Exception as e_fallback: - logger.error(f"Error with LiteLLM API fallback: {str(e_fallback)}") - error_msg = f"\nError with LiteLLM API fallback: {str(e_fallback)}\n\nPlease check that you have set the LITELLM_API_KEY environment variable with a valid API key." - await websocket.send_text(error_msg) - elif request.provider == "bedrock": - try: - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM, - ) - - logger.info("Making fallback AWS Bedrock API call") - fallback_response = await model.acall( - api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM - ) - - if isinstance(fallback_response, str): - await websocket.send_text(fallback_response) - else: - await websocket.send_text(str(fallback_response)) - except Exception as e_fallback: - logger.error( - f"Error with AWS Bedrock API fallback: {str(e_fallback)}" - ) - error_msg = ( - f"\nError with AWS Bedrock API fallback: {str(e_fallback)}\n\n" - "Please check that you have set the AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY " - "environment variables with valid credentials." - ) - await websocket.send_text(error_msg) - elif request.provider == "azure": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM - ) - - # Get the response using the simplified prompt - logger.info("Making fallback Azure AI API call") - fallback_response = await model.acall(api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM) - - # Handle streaming fallback response from Azure AI - async for chunk in fallback_response: - choices = getattr(chunk, "choices", []) - if len(choices) > 0: - delta = getattr(choices[0], "delta", None) - if delta is not None: - text = getattr(delta, "content", None) - if text is not None: - await websocket.send_text(text) - except Exception as e_fallback: - logger.error(f"Error with Azure AI API fallback: {str(e_fallback)}") - error_msg = f"\nError with Azure AI API fallback: {str(e_fallback)}\n\nPlease check that you have set the AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, and AZURE_OPENAI_VERSION environment variables with valid values." - await websocket.send_text(error_msg) - elif request.provider == "dashscope": - try: - # Create new api_kwargs with the simplified prompt - fallback_api_kwargs = model.convert_inputs_to_api_kwargs( - input=simplified_prompt, - model_kwargs=model_kwargs, - model_type=ModelType.LLM, - ) - - logger.info("Making fallback Dashscope API call") - fallback_response = await model.acall( - api_kwargs=fallback_api_kwargs, model_type=ModelType.LLM - ) - - # DashscopeClient.acall (stream=True) returns an async - # generator of text chunks - async for text in fallback_response: - if text: - await websocket.send_text(text) - except Exception as e_fallback: - logger.error( - f"Error with Dashscope API fallback: {str(e_fallback)}" - ) - error_msg = ( - f"\nError with Dashscope API fallback: {str(e_fallback)}\n\n" - "Please check that you have set the DASHSCOPE_API_KEY (and optionally " - "DASHSCOPE_WORKSPACE_ID) environment variables with valid values." - ) - await websocket.send_text(error_msg) - else: - # Google Generative AI fallback (default provider) - model_config = get_model_config(request.provider, request.model) - fallback_model = genai.GenerativeModel( - model_name=model_config["model_kwargs"]["model"], - generation_config={ - "temperature": model_config["model_kwargs"].get("temperature", 0.7), - "top_p": model_config["model_kwargs"].get("top_p", 0.8), - "top_k": model_config["model_kwargs"].get("top_k", 40), - }, + async def stream_and_fallback( + streamer: ChatStreamer, + prompt_func: Callable[[], str], + simplified_prompt_func: Callable[[], str], + ) -> AsyncIterator[str]: + try: + async for chunk in streamer.respond_stream(prompt_func()): + yield chunk + except Exception as e: + if is_token_limit_error(e): + logger.warning("Token limit exceeded, retrying without context") + try: + async for chunk in streamer.respond_stream(simplified_prompt_func()): + yield chunk + except Exception as e2: + logger.error("Error in fallback streaming response: %s", e2) + yield ( + "\nI apologize, but your request is too large for me to process. " + "Please try a shorter query or break it into smaller parts." ) + else: + msg = f"Error with {streamer.provider} API: {e}" + logger.error(msg, exc_info=True) + if streamer.error_hint: + msg += f"\n\n{streamer.error_hint}" + yield "\n" + msg - fallback_response = fallback_model.generate_content( - simplified_prompt, stream=True - ) - for chunk in fallback_response: - if hasattr(chunk, "text"): - await websocket.send_text(chunk.text) - except Exception as e2: - logger.error(f"Error in fallback streaming response: {str(e2)}") - await websocket.send_text(f"\nI apologize, but your request is too large for me to process. Please try a shorter query or break it into smaller parts.") - # Close the WebSocket connection after sending the error message - await websocket.close() - else: - # For other errors, return the error message - await websocket.send_text(f"\nError: {error_message}") - # Close the WebSocket connection after sending the error message - await websocket.close() + model_config = get_model_config(request.provider, request.model)["model_kwargs"] + chat_streamer = ChatStreamer.create( + provider=request.provider, + model=request.model, + model_config=model_config, + ) + + prompt_kwargs = dict( + system_prompt=system_prompt, + query=query, + conversation_history=conversation_history, + file_path=request.filePath, + file_content=file_content, + context=context_text, + ) + prompt_func = partial(prompt_builder, **prompt_kwargs, simplify=False) + simplified_prompt_func = partial(prompt_builder, **prompt_kwargs, simplify=True) + + async for chunk in stream_and_fallback(chat_streamer, prompt_func, simplified_prompt_func): + await websocket.send_text(chunk) + await websocket.close() except WebSocketDisconnect: logger.info("WebSocket disconnected") diff --git a/tests/unit/test_chat.py b/tests/unit/test_chat.py index 181791626..5089189ee 100644 --- a/tests/unit/test_chat.py +++ b/tests/unit/test_chat.py @@ -1,6 +1,6 @@ import pytest -from api.chat import ( - ChatStreamer, +from api.chat import ChatStreamer +from api.chat._stream import ( OllamaChatStreamer, OpenRouterChatStreamer, OpenAIChatStreamer, @@ -8,6 +8,7 @@ BedrockChatStreamer, DashScopeChatStreamer, GoogleGenerativeChatStreamer, + LiteLLMChatStreamer, ) @pytest.mark.parametrize("provider, expected", [ @@ -18,6 +19,7 @@ ("bedrock", BedrockChatStreamer), ("dashscope", DashScopeChatStreamer), ("google", GoogleGenerativeChatStreamer), + ("litellm", LiteLLMChatStreamer), ]) def test_every_provider_is_registered(provider, expected): assert ChatStreamer._registry[provider] is expected @@ -30,6 +32,7 @@ def test_every_provider_is_registered(provider, expected): ("bedrock", BedrockChatStreamer), ("dashscope", DashScopeChatStreamer), ("google", GoogleGenerativeChatStreamer), + ("litellm", LiteLLMChatStreamer), ]) def test_create_returns_correct_subclass(monkeypatch, provider, expected): monkeypatch.setattr(expected, "__init__", lambda self, **kw: None) From a601f2b2569e7db3e1e733cd26196271f6df0df4 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 18 Jul 2026 21:20:39 +0800 Subject: [PATCH 07/25] fix: fix simple_chat format --- api/simple_chat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/simple_chat.py b/api/simple_chat.py index ab1d9de39..9ac39c43c 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -8,7 +8,7 @@ from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field -from api.chat import ChatStreamer, , prompt_builder, is_token_limit_error +from api.chat import ChatStreamer, prompt_builder, is_token_limit_error from api.config import get_model_config, configs from api.data_pipeline import count_tokens, get_file_content from api.rag import RAG, MAX_INPUT_TOKENS From 5066942fd841932a2c378a6e98e16a5d140cf0ae Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 18 Jul 2026 23:26:15 +0800 Subject: [PATCH 08/25] refactor and add clients module (#7) --- api/azureai_client.py | 488 -------------- api/chat.py | 279 -------- api/chat/_stream.py | 14 +- api/clients/__init__.py | 19 + api/{bedrock_client.py => clients/bedrock.py} | 0 .../dashscope.py} | 0 .../google_embedder.py} | 0 api/{litellm_client.py => clients/litellm.py} | 2 +- api/{ollama_client.py => clients/ollama.py} | 0 .../openrouter.py} | 0 api/config.py | 20 +- api/openai_client.py | 629 ------------------ 12 files changed, 38 insertions(+), 1413 deletions(-) delete mode 100644 api/azureai_client.py delete mode 100644 api/chat.py create mode 100644 api/clients/__init__.py rename api/{bedrock_client.py => clients/bedrock.py} (100%) rename api/{dashscope_client.py => clients/dashscope.py} (100%) rename api/{google_embedder_client.py => clients/google_embedder.py} (100%) rename api/{litellm_client.py => clients/litellm.py} (96%) rename api/{ollama_client.py => clients/ollama.py} (100%) rename api/{openrouter_client.py => clients/openrouter.py} (100%) delete mode 100644 api/openai_client.py diff --git a/api/azureai_client.py b/api/azureai_client.py deleted file mode 100644 index 948e86c30..000000000 --- a/api/azureai_client.py +++ /dev/null @@ -1,488 +0,0 @@ -"""AzureOpenAI ModelClient integration.""" - -import os -from typing import ( - Dict, - Sequence, - Optional, - List, - Any, - TypeVar, - Callable, - Generator, - Union, - Literal, -) -import re - -import logging -import backoff - -# optional import -from adalflow.utils.lazy_import import safe_import, OptionalPackages - -import sys - -openai = safe_import(OptionalPackages.OPENAI.value[0], OptionalPackages.OPENAI.value[1]) -# Importing all Azure packages together -azure_modules = safe_import( - OptionalPackages.AZURE.value[0], # List of package names - OptionalPackages.AZURE.value[1], # Error message -) -# Manually add each module to sys.modules to make them available globally as if imported normally -azure_module_names = OptionalPackages.AZURE.value[0] -for name, module in zip(azure_module_names, azure_modules): - sys.modules[name] = module - -# Use the modules as if they were imported normally -from azure.identity import DefaultAzureCredential, get_bearer_token_provider - -# from azure.core.credentials import AccessToken -from openai import AzureOpenAI, AsyncAzureOpenAI, Stream -from openai import ( - APITimeoutError, - InternalServerError, - RateLimitError, - UnprocessableEntityError, - BadRequestError, -) -from openai.types import ( - Completion, - CreateEmbeddingResponse, -) -from openai.types.chat import ChatCompletionChunk, ChatCompletion - -from adalflow.core.model_client import ModelClient -from adalflow.core.types import ( - ModelType, - EmbedderOutput, - TokenLogProb, - CompletionUsage, - GeneratorOutput, -) -from adalflow.components.model_client.utils import parse_embedding_response - -log = logging.getLogger(__name__) -T = TypeVar("T") - - -__all__ = ["AzureAIClient"] - -# TODO: this overlaps with openai client largely, might need to refactor to subclass openai client to simplify the code - - -# completion parsing functions and you can combine them into one singple chat completion parser -def get_first_message_content(completion: ChatCompletion) -> str: - r"""When we only need the content of the first message. - It is the default parser for chat completion.""" - return completion.choices[0].message.content - - -# def _get_chat_completion_usage(completion: ChatCompletion) -> OpenAICompletionUsage: -# return completion.usage - - -def parse_stream_response(completion: ChatCompletionChunk) -> str: - r"""Parse the response of the stream API.""" - return completion.choices[0].delta.content - - -def handle_streaming_response(generator: Stream[ChatCompletionChunk]): - r"""Handle the streaming response.""" - for completion in generator: - log.debug(f"Raw chunk completion: {completion}") - parsed_content = parse_stream_response(completion) - yield parsed_content - - -def get_all_messages_content(completion: ChatCompletion) -> List[str]: - r"""When the n > 1, get all the messages content.""" - return [c.message.content for c in completion.choices] - - -def get_probabilities(completion: ChatCompletion) -> List[List[TokenLogProb]]: - r"""Get the probabilities of each token in the completion.""" - log_probs = [] - for c in completion.choices: - content = c.logprobs.content - print(content) - log_probs_for_choice = [] - for openai_token_logprob in content: - token = openai_token_logprob.token - logprob = openai_token_logprob.logprob - log_probs_for_choice.append(TokenLogProb(token=token, logprob=logprob)) - log_probs.append(log_probs_for_choice) - return log_probs - - -class AzureAIClient(ModelClient): - __doc__ = r""" - A client wrapper for interacting with Azure OpenAI's API. - - This class provides support for both embedding and chat completion API calls. - Users can use this class to simplify their interactions with Azure OpenAI models - through the `Embedder` and `Generator` components. - - **Initialization:** - - You can initialize the `AzureAIClient` with either an API key or Azure Active Directory (AAD) token - authentication. It is recommended to set environment variables for sensitive data like API keys. - - Args: - api_key (Optional[str]): Azure OpenAI API key. Default is None. - api_version (Optional[str]): API version to use. Default is None. - azure_endpoint (Optional[str]): Azure OpenAI endpoint URL. Default is None. - credential (Optional[DefaultAzureCredential]): Azure AD credential for token-based authentication. Default is None. - chat_completion_parser (Callable[[Completion], Any]): Function to parse chat completions. Default is `get_first_message_content`. - input_type (Literal["text", "messages"]): Format for input, either "text" or "messages". Default is "text". - - **Setup Instructions:** - - - **Using API Key:** - Set up the following environment variables: - ```bash - export AZURE_OPENAI_API_KEY="your_api_key" - export AZURE_OPENAI_ENDPOINT="your_endpoint" - export AZURE_OPENAI_VERSION="your_version" - ``` - - - **Using Azure AD Token:** - Ensure you have configured Azure AD credentials. The `DefaultAzureCredential` will automatically use your configured credentials. - - **Example Usage:** - - .. code-block:: python - - from azure.identity import DefaultAzureCredential - from your_module import AzureAIClient # Adjust import based on your module name - - # Initialize with API key - client = AzureAIClient( - api_key="your_api_key", - api_version="2023-05-15", - azure_endpoint="https://your-endpoint.openai.azure.com/" - ) - - # Or initialize with Azure AD token - client = AzureAIClient( - api_version="2023-05-15", - azure_endpoint="https://your-endpoint.openai.azure.com/", - credential=DefaultAzureCredential() - ) - - # Example call to the chat completion API - api_kwargs = { - "model": "gpt-3.5-turbo", - "messages": [{"role": "user", "content": "What is the meaning of life?"}], - "stream": True - } - response = client.call(api_kwargs=api_kwargs, model_type=ModelType.LLM) - - for chunk in response: - print(chunk) - - - **Notes:** - - Ensure that the API key or credentials are correctly set up and accessible to avoid authentication errors. - - Use `chat_completion_parser` to define how to extract and handle the chat completion responses. - - The `input_type` parameter determines how input is formatted for the API call. - - **References:** - - [Azure OpenAI API Documentation](https://learn.microsoft.com/en-us/azure/ai-services/openai/overview) - - [OpenAI API Documentation](https://platform.openai.com/docs/guides/text-generation) - """ - - def __init__( - self, - api_key: Optional[str] = None, - api_version: Optional[str] = None, - azure_endpoint: Optional[str] = None, - credential: Optional[DefaultAzureCredential] = None, - chat_completion_parser: Callable[[Completion], Any] = None, - input_type: Literal["text", "messages"] = "text", - ): - r"""It is recommended to set the API_KEY into the environment variable instead of passing it as an argument. - - - Initializes the Azure OpenAI client with either API key or AAD token authentication. - - Args: - api_key: Azure OpenAI API key. - api_version: Azure OpenAI API version. - azure_endpoint: Azure OpenAI endpoint. - credential: Azure AD credential for token-based authentication. - chat_completion_parser: Function to parse chat completions. - input_type: Input format, either "text" or "messages". - - """ - super().__init__() - - # added api_type azure for azure Ai - self.api_type = "azure" - self._api_key = api_key - self._apiversion = api_version - self._azure_endpoint = azure_endpoint - self._credential = credential - self.sync_client = self.init_sync_client() - self.async_client = None # only initialize if the async call is called - self.chat_completion_parser = ( - chat_completion_parser or get_first_message_content - ) - self._input_type = input_type - - def init_sync_client(self): - api_key = self._api_key or os.getenv("AZURE_OPENAI_API_KEY") - azure_endpoint = self._azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT") - api_version = self._apiversion or os.getenv("AZURE_OPENAI_VERSION") - # credential = self._credential or DefaultAzureCredential - if not azure_endpoint: - raise ValueError("Environment variable AZURE_OPENAI_ENDPOINT must be set") - if not api_version: - raise ValueError("Environment variable AZURE_OPENAI_VERSION must be set") - - if api_key: - return AzureOpenAI( - api_key=api_key, azure_endpoint=azure_endpoint, api_version=api_version - ) - elif self._credential: - # credential = DefaultAzureCredential() - token_provider = get_bearer_token_provider( - DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default" - ) - return AzureOpenAI( - azure_ad_token_provider=token_provider, - azure_endpoint=azure_endpoint, - api_version=api_version, - ) - else: - raise ValueError( - "Environment variable AZURE_OPENAI_API_KEY must be set or credential must be provided" - ) - - def init_async_client(self): - api_key = self._api_key or os.getenv("AZURE_OPENAI_API_KEY") - azure_endpoint = self._azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT") - api_version = self._apiversion or os.getenv("AZURE_OPENAI_VERSION") - # credential = self._credential or DefaultAzureCredential() - if not azure_endpoint: - raise ValueError("Environment variable AZURE_OPENAI_ENDPOINT must be set") - if not api_version: - raise ValueError("Environment variable AZURE_OPENAI_VERSION must be set") - - if api_key: - return AsyncAzureOpenAI( - api_key=api_key, azure_endpoint=azure_endpoint, api_version=api_version - ) - elif self._credential: - # credential = DefaultAzureCredential() - token_provider = get_bearer_token_provider( - DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default" - ) - return AsyncAzureOpenAI( - azure_ad_token_provider=token_provider, - azure_endpoint=azure_endpoint, - api_version=api_version, - ) - else: - raise ValueError( - "Environment variable AZURE_OPENAI_API_KEY must be set or credential must be provided" - ) - - # def _parse_chat_completion(self, completion: ChatCompletion) -> "GeneratorOutput": - # # TODO: raw output it is better to save the whole completion as a source of truth instead of just the message - # try: - # data = self.chat_completion_parser(completion) - # usage = self.track_completion_usage(completion) - # return GeneratorOutput( - # data=data, error=None, raw_response=str(data), usage=usage - # ) - # except Exception as e: - # log.error(f"Error parsing the completion: {e}") - # return GeneratorOutput(data=None, error=str(e), raw_response=completion) - - def parse_chat_completion( - self, - completion: Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]], - ) -> "GeneratorOutput": - """Parse the completion, and put it into the raw_response.""" - log.debug(f"completion: {completion}, parser: {self.chat_completion_parser}") - try: - data = self.chat_completion_parser(completion) - usage = self.track_completion_usage(completion) - return GeneratorOutput( - data=None, error=None, raw_response=data, usage=usage - ) - except Exception as e: - log.error(f"Error parsing the completion: {e}") - return GeneratorOutput(data=None, error=str(e), raw_response=completion) - - def track_completion_usage( - self, - completion: Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]], - ) -> CompletionUsage: - if isinstance(completion, ChatCompletion): - usage: CompletionUsage = CompletionUsage( - completion_tokens=completion.usage.completion_tokens, - prompt_tokens=completion.usage.prompt_tokens, - total_tokens=completion.usage.total_tokens, - ) - return usage - else: - raise NotImplementedError( - "streaming completion usage tracking is not implemented" - ) - - def parse_embedding_response( - self, response: CreateEmbeddingResponse - ) -> EmbedderOutput: - r"""Parse the embedding response to a structure AdalFlow components can understand. - - Should be called in ``Embedder``. - """ - try: - return parse_embedding_response(response) - except Exception as e: - log.error(f"Error parsing the embedding response: {e}") - return EmbedderOutput(data=[], error=str(e), raw_response=response) - - def convert_inputs_to_api_kwargs( - self, - input: Optional[Any] = None, - model_kwargs: Dict = {}, - model_type: ModelType = ModelType.UNDEFINED, - ) -> Dict: - r""" - Specify the API input type and output api_kwargs that will be used in _call and _acall methods. - Convert the Component's standard input, and system_input(chat model) and model_kwargs into API-specific format - """ - - final_model_kwargs = model_kwargs.copy() - if model_type == ModelType.EMBEDDER: - if isinstance(input, str): - input = [input] - # convert input to input - if not isinstance(input, Sequence): - raise TypeError("input must be a sequence of text") - final_model_kwargs["input"] = input - elif model_type == ModelType.LLM: - # convert input to messages - messages: List[Dict[str, str]] = [] - - if self._input_type == "messages": - system_start_tag = "" - system_end_tag = "" - user_start_tag = "" - user_end_tag = "" - pattern = f"{system_start_tag}(.*?){system_end_tag}{user_start_tag}(.*?){user_end_tag}" - # Compile the regular expression - regex = re.compile(pattern) - # Match the pattern - match = regex.search(input) - system_prompt, input_str = None, None - - if match: - system_prompt = match.group(1) - input_str = match.group(2) - - else: - print("No match found.") - if system_prompt and input_str: - messages.append({"role": "system", "content": system_prompt}) - messages.append({"role": "user", "content": input_str}) - if len(messages) == 0: - messages.append({"role": "system", "content": input}) - final_model_kwargs["messages"] = messages - else: - raise ValueError(f"model_type {model_type} is not supported") - return final_model_kwargs - - @backoff.on_exception( - backoff.expo, - ( - APITimeoutError, - InternalServerError, - RateLimitError, - UnprocessableEntityError, - BadRequestError, - ), - max_time=5, - ) - def call(self, api_kwargs: Dict = {}, model_type: ModelType = ModelType.UNDEFINED): - """ - kwargs is the combined input and model_kwargs. Support streaming call. - """ - log.info(f"api_kwargs: {api_kwargs}") - if model_type == ModelType.EMBEDDER: - return self.sync_client.embeddings.create(**api_kwargs) - elif model_type == ModelType.LLM: - if "stream" in api_kwargs and api_kwargs.get("stream", False): - log.debug("streaming call") - self.chat_completion_parser = handle_streaming_response - return self.sync_client.chat.completions.create(**api_kwargs) - return self.sync_client.chat.completions.create(**api_kwargs) - else: - raise ValueError(f"model_type {model_type} is not supported") - - @backoff.on_exception( - backoff.expo, - ( - APITimeoutError, - InternalServerError, - RateLimitError, - UnprocessableEntityError, - BadRequestError, - ), - max_time=5, - ) - async def acall( - self, api_kwargs: Dict = {}, model_type: ModelType = ModelType.UNDEFINED - ): - """ - kwargs is the combined input and model_kwargs - """ - if self.async_client is None: - self.async_client = self.init_async_client() - if model_type == ModelType.EMBEDDER: - return await self.async_client.embeddings.create(**api_kwargs) - elif model_type == ModelType.LLM: - return await self.async_client.chat.completions.create(**api_kwargs) - else: - raise ValueError(f"model_type {model_type} is not supported") - - @classmethod - def from_dict(cls: type[T], data: Dict[str, Any]) -> T: - obj = super().from_dict(data) - # recreate the existing clients - obj.sync_client = obj.init_sync_client() - obj.async_client = obj.init_async_client() - return obj - - def to_dict(self) -> Dict[str, Any]: - r"""Convert the component to a dictionary.""" - # TODO: not exclude but save yes or no for recreating the clients - exclude = [ - "sync_client", - "async_client", - ] # unserializable object - output = super().to_dict(exclude=exclude) - return output - - -# if __name__ == "__main__": -# from adalflow.core import Generator -# from adalflow.utils import setup_env, get_logger - -# log = get_logger(level="DEBUG") - -# setup_env() -# prompt_kwargs = {"input_str": "What is the meaning of life?"} - -# gen = Generator( -# model_client=OpenAIClient(), -# model_kwargs={"model": "gpt-3.5-turbo", "stream": True}, -# ) -# gen_response = gen(prompt_kwargs) -# print(f"gen_response: {gen_response}") - -# for genout in gen_response.data: -# print(f"genout: {genout}") \ No newline at end of file diff --git a/api/chat.py b/api/chat.py deleted file mode 100644 index 60b05aba8..000000000 --- a/api/chat.py +++ /dev/null @@ -1,279 +0,0 @@ -import logging -from abc import abstractmethod, ABC - -from typing import TYPE_CHECKING - -from collections.abc import AsyncIterator - -from adalflow.core.types import ModelType -from api.config import ( - OPENROUTER_API_KEY, - OPENAI_API_KEY, - AWS_ACCESS_KEY_ID, - AWS_SECRET_ACCESS_KEY, -) - -if TYPE_CHECKING: - from ollama import ChatResponse - from openai.types.chat import ChatCompletionChunk - from openai import AsyncStream - -MODEL_CFG = dict[str, str | int | float] - -# Configure logging -from api.logging_config import setup_logging - -setup_logging() -logger = logging.getLogger(__name__) - - -class ChatStreamer(ABC): - _registry: dict[str, type["ChatStreamer"]] = {} - provider: str - - def __init_subclass__(cls, **kwargs) -> None: - super().__init_subclass__(**kwargs) - if provider := getattr(cls, "provider", None): - ChatStreamer._registry[provider] = cls - - @classmethod - def create(cls, *, provider: str, model: str | None = None, model_config: MODEL_CFG) -> "ChatStreamer": - model = model or model_config.get("model") - logger.info("Using %s with model: %s", provider, model) - registered = ChatStreamer._registry.get(provider, None) - if registered: - return registered(model=model, model_config=model_config) - raise RuntimeError(f"Provider {provider} not registered") - - @abstractmethod - def respond_stream(self, prompt: str) -> AsyncIterator[str]: - raise NotImplementedError(f"{type(self).__name__} does not implement `respond_stream`") - - -class OllamaChatStreamer(ChatStreamer): - provider = "ollama" - - def __init__(self, *, model: str, model_config: MODEL_CFG): - from adalflow.components.model_client.ollama_client import OllamaClient - - self.client = OllamaClient() - self.model_kwargs = { - "model": model, - "stream": True, - "options": { - "temperature": model_config["temperature"], - "top_p": model_config["top_p"], - "num_ctx": model_config["num_ctx"] - } - } - - logger.debug(f"Prompting Ollama with kwargs: {self.model_kwargs}") - - async def respond_stream(self, prompt: str) -> AsyncIterator[str]: - api_kwargs = self.client.convert_inputs_to_api_kwargs( - input=prompt + " /no_think", # todo I think this could be added into model kwargs? - model_kwargs=self.model_kwargs, - model_type=ModelType.LLM, - ) - - response: "AsyncIterator[ChatResponse]" = await self.client.acall( - api_kwargs=api_kwargs, - model_type=ModelType.LLM, - ) - - async for chunk in response: - if not hasattr(chunk, "message"): - raise RuntimeError( - "`message` field not found in response. Wrong ollama-python version probably.", - ) - text = chunk.message.content - if text: - text = text.replace('', '').replace('', '') - yield text - - -class OpenRouterChatStreamer(ChatStreamer): - provider = "openrouter" - - def __init__(self, *, model: str, model_config: MODEL_CFG): - if not OPENROUTER_API_KEY: - logger.warning("OPENROUTER_API_KEY not configured, but continuing with request") - # We'll let the OpenRouterClient handle this and return a friendly error message - from api.openrouter_client import OpenRouterClient - - self.client = OpenRouterClient() - self.model_kwargs = { - "model": model, - "stream": True, - "temperature": model_config["temperature"] - } - if "top_k" in model_config: - self.model_kwargs["top_k"] = model_config["top_k"] - - async def respond_stream(self, prompt: str) -> AsyncIterator[str]: - api_kwargs = self.client.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=self.model_kwargs, - model_type=ModelType.LLM, - ) - async for chunk in await self.client.acall( - api_kwargs=api_kwargs, - model_type=ModelType.LLM, - ): - yield chunk - - -class OpenAIChatStreamer(ChatStreamer): - provider = "openai" - - def __init__(self, *, model: str, model_config: MODEL_CFG): - if not OPENAI_API_KEY: - logger.warning("OPENAI_API_KEY not configured, but continuing with request") - # We'll let the OpenAIClient handle this and return an error message - from api.openai_client import OpenAIClient - - self.client = OpenAIClient() - self.model_kwargs = { - "model": model, - "stream": True, - "temperature": model_config["temperature"] - } - # Only add top_p if it exists in the model config - if "top_p" in model_config: - self.model_kwargs["top_p"] = model_config["top_p"] - - async def respond_stream(self, prompt: str) -> AsyncIterator[str]: - api_kwargs = self.client.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=self.model_kwargs, - model_type=ModelType.LLM - ) - response: "AsyncStream[ChatCompletionChunk]" = await self.client.acall( - api_kwargs=api_kwargs, - model_type=ModelType.LLM, - ) - - async for chunk in response: - if ( - chunk.choices and - chunk.choices[0].delta is not None and - chunk.choices[0].delta.content is not None - ): - yield chunk.choices[0].delta.content - - -class AzureChatStreamer(ChatStreamer): - provider = "azure" - - def __init__(self, *, model: str, model_config: MODEL_CFG): - from api.azureai_client import AzureAIClient - - self.client = AzureAIClient() - self.model_kwargs = { - "model": model, - "stream": True, - "temperature": model_config["temperature"], - "top_p": model_config["top_p"] - } - - async def respond_stream(self, prompt: str) -> AsyncIterator[str]: - api_kwargs = self.client.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=self.model_kwargs, - model_type=ModelType.LLM - ) - response: "AsyncStream[ChatCompletionChunk]" = await self.client.acall( - api_kwargs=api_kwargs, - model_type=ModelType.LLM, - ) - - async for chunk in response: - if ( - chunk.choices and - chunk.choices[0].delta is not None and - chunk.choices[0].delta.content is not None - ): - yield chunk.choices[0].delta.content - - -class BedrockChatStreamer(ChatStreamer): - provider = "bedrock" - - def __init__(self, *, model: str, model_config: MODEL_CFG): - if not AWS_ACCESS_KEY_ID or not AWS_SECRET_ACCESS_KEY: - logger.warning("AWS_ACCESS_KEY_ID or AWS_SECRET_ACCESS_KEY not configured, but continuing with request") - # We'll let the BedrockClient handle this and return an error message - from api.bedrock_client import BedrockClient - - self.client = BedrockClient() - self.model_kwargs = { - "model": model, - "temperature": model_config["temperature"], - "top_p": model_config["top_p"] - } - - async def respond_stream(self, prompt: str) -> AsyncIterator[str]: - api_kwargs = self.client.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=self.model_kwargs, - model_type=ModelType.LLM, - ) - response = await self.client.acall( - api_kwargs=api_kwargs, - model_type=ModelType.LLM, - ) - if not isinstance(response, str): - response = str(response) - yield response - - -class DashScopeChatStreamer(ChatStreamer): - provider = "dashscope" - - def __init__(self, *, model: str, model_config: MODEL_CFG): - from api.dashscope_client import DashscopeClient - - self.client = DashscopeClient() - self.model_kwargs = { - "model": model, - "stream": True, - "temperature": model_config["temperature"], - "top_p": model_config["top_p"], - } - - async def respond_stream(self, prompt: str) -> AsyncIterator[str]: - api_kwargs = self.client.convert_inputs_to_api_kwargs( - input=prompt, - model_kwargs=self.model_kwargs, - model_type=ModelType.LLM, - ) - response = await self.client.acall( - api_kwargs=api_kwargs, - model_type=ModelType.LLM, - ) - async for text in response: - if text: - yield text - - -class GoogleGenerativeChatStreamer(ChatStreamer): - provider = "google" - - def __init__(self, *, model: str, model_config: MODEL_CFG): - import google.generativeai as genai - from google.generativeai.types import GenerationConfig - - self.client = genai.GenerativeModel( - model_name=model, - generation_config=GenerationConfig( - temperature=model_config["temperature"], - top_p=model_config["top_p"], - top_k=model_config["top_k"], - ) - ) - - async def respond_stream(self, prompt: str) -> AsyncIterator[str]: - response = await self.client.generate_content_async(prompt, stream=True) - async for chunk in response: - if hasattr(chunk, "text"): - yield chunk.text diff --git a/api/chat/_stream.py b/api/chat/_stream.py index 94251d162..5e11bb871 100644 --- a/api/chat/_stream.py +++ b/api/chat/_stream.py @@ -18,7 +18,7 @@ from ollama import ChatResponse from openai.types.chat import ChatCompletionChunk from openai import AsyncStream - from api.openai_client import OpenAIClient + from api.clients import OpenAIClient MODEL_CFG = dict[str, str | int | float] @@ -106,7 +106,7 @@ def __init__(self, *, model: str, model_config: MODEL_CFG): if not OPENROUTER_API_KEY: logger.warning("OPENROUTER_API_KEY not configured, but continuing with request") # We'll let the OpenRouterClient handle this and return a friendly error message - from api.openrouter_client import OpenRouterClient + from api.clients import OpenRouterClient self.client = OpenRouterClient() self.model_kwargs = { @@ -185,7 +185,7 @@ def __init__(self, *, model: str, model_config: MODEL_CFG): super().__init__(model=model, model_config=model_config) def _build_client(self): - from api.openai_client import OpenAIClient + from api.clients import OpenAIClient return OpenAIClient() @@ -198,7 +198,7 @@ class AzureChatStreamer(_OpenAICompatStreamer): ) def _build_client(self): - from api.azureai_client import AzureAIClient + from adalflow.components.model_client import AzureAIClient return AzureAIClient() @@ -217,7 +217,7 @@ def __init__(self, *, model: str, model_config: MODEL_CFG): super().__init__(model=model, model_config=model_config) def _build_client(self): - from api.litellm_client import LiteLLMClient + from api.clients import LiteLLMClient return LiteLLMClient() @@ -232,7 +232,7 @@ def __init__(self, *, model: str, model_config: MODEL_CFG): if not AWS_ACCESS_KEY_ID or not AWS_SECRET_ACCESS_KEY: logger.warning("AWS_ACCESS_KEY_ID or AWS_SECRET_ACCESS_KEY not configured, but continuing with request") # We'll let the BedrockClient handle this and return an error message - from api.bedrock_client import BedrockClient + from api.clients import BedrockClient self.client = BedrockClient() self.model_kwargs = {"model": model} @@ -267,7 +267,7 @@ class DashScopeChatStreamer(ChatStreamer): ) def __init__(self, *, model: str, model_config: MODEL_CFG): - from api.dashscope_client import DashscopeClient + from api.clients import DashscopeClient self.client = DashscopeClient() self.model_kwargs = { diff --git a/api/clients/__init__.py b/api/clients/__init__.py new file mode 100644 index 000000000..720a1e4d8 --- /dev/null +++ b/api/clients/__init__.py @@ -0,0 +1,19 @@ +from .google_embedder import GoogleEmbedderClient +from .bedrock import BedrockClient +from .dashscope import DashscopeClient +from .ollama import OllamaClient +from .litellm import LiteLLMClient +from .openrouter import OpenRouterClient +from adalflow.components.model_client import AzureAIClient, OpenAIClient, GoogleGenAIClient + +__all__ = [ + "AzureAIClient", + "BedrockClient", + "DashscopeClient", + "GoogleEmbedderClient", + "GoogleGenAIClient", + "LiteLLMClient", + "OllamaClient", + "OpenAIClient", + "OpenRouterClient", +] \ No newline at end of file diff --git a/api/bedrock_client.py b/api/clients/bedrock.py similarity index 100% rename from api/bedrock_client.py rename to api/clients/bedrock.py diff --git a/api/dashscope_client.py b/api/clients/dashscope.py similarity index 100% rename from api/dashscope_client.py rename to api/clients/dashscope.py diff --git a/api/google_embedder_client.py b/api/clients/google_embedder.py similarity index 100% rename from api/google_embedder_client.py rename to api/clients/google_embedder.py diff --git a/api/litellm_client.py b/api/clients/litellm.py similarity index 96% rename from api/litellm_client.py rename to api/clients/litellm.py index 439adced6..1c814d11f 100644 --- a/api/litellm_client.py +++ b/api/clients/litellm.py @@ -2,7 +2,7 @@ from typing import Optional, Callable from openai import AsyncOpenAI, OpenAI -from api.openai_client import OpenAIClient +from adalflow.components.model_client.openai_client import OpenAIClient class LiteLLMClient(OpenAIClient): diff --git a/api/ollama_client.py b/api/clients/ollama.py similarity index 100% rename from api/ollama_client.py rename to api/clients/ollama.py diff --git a/api/openrouter_client.py b/api/clients/openrouter.py similarity index 100% rename from api/openrouter_client.py rename to api/clients/openrouter.py diff --git a/api/config.py b/api/config.py index 4cafcdbc7..be11c368a 100644 --- a/api/config.py +++ b/api/config.py @@ -7,15 +7,17 @@ logger = logging.getLogger(__name__) -from api.openai_client import OpenAIClient -from api.litellm_client import LiteLLMClient -from api.openrouter_client import OpenRouterClient -from api.bedrock_client import BedrockClient -from api.google_embedder_client import GoogleEmbedderClient -from api.azureai_client import AzureAIClient -from api.dashscope_client import DashscopeClient -from api.ollama_client import OllamaClient -from adalflow import GoogleGenAIClient +from api.clients import ( + LiteLLMClient, + OpenRouterClient, + BedrockClient, + GoogleGenAIClient, + GoogleEmbedderClient, + AzureAIClient, + DashscopeClient, + OpenAIClient, + OllamaClient, +) # Get API keys from environment variables OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') diff --git a/api/openai_client.py b/api/openai_client.py deleted file mode 100644 index bc75ed586..000000000 --- a/api/openai_client.py +++ /dev/null @@ -1,629 +0,0 @@ -"""OpenAI ModelClient integration.""" - -import os -import base64 -from typing import ( - Dict, - Sequence, - Optional, - List, - Any, - TypeVar, - Callable, - Generator, - Union, - Literal, -) -import re - -import logging -import backoff - -# optional import -from adalflow.utils.lazy_import import safe_import, OptionalPackages -from openai.types.chat.chat_completion import Choice - -openai = safe_import(OptionalPackages.OPENAI.value[0], OptionalPackages.OPENAI.value[1]) - -from openai import OpenAI, AsyncOpenAI, Stream -from openai import ( - APITimeoutError, - InternalServerError, - RateLimitError, - UnprocessableEntityError, - BadRequestError, -) -from openai.types import ( - Completion, - CreateEmbeddingResponse, - Image, -) -from openai.types.chat import ChatCompletionChunk, ChatCompletion, ChatCompletionMessage - -from adalflow.core.model_client import ModelClient -from adalflow.core.types import ( - ModelType, - EmbedderOutput, - TokenLogProb, - CompletionUsage, - GeneratorOutput, -) -from adalflow.components.model_client.utils import parse_embedding_response - -log = logging.getLogger(__name__) -T = TypeVar("T") - - -# completion parsing functions and you can combine them into one singple chat completion parser -def get_first_message_content(completion: ChatCompletion) -> str: - r"""When we only need the content of the first message. - It is the default parser for chat completion.""" - log.debug(f"raw completion: {completion}") - return completion.choices[0].message.content - - -# def _get_chat_completion_usage(completion: ChatCompletion) -> OpenAICompletionUsage: -# return completion.usage - - -# A simple heuristic to estimate token count for estimating number of tokens in a Streaming response -def estimate_token_count(text: str) -> int: - """ - Estimate the token count of a given text. - - Args: - text (str): The text to estimate token count for. - - Returns: - int: Estimated token count. - """ - # Split the text into tokens using spaces as a simple heuristic - tokens = text.split() - - # Return the number of tokens - return len(tokens) - - -def parse_stream_response(completion: ChatCompletionChunk) -> str: - r"""Parse the response of the stream API.""" - return completion.choices[0].delta.content - - -def handle_streaming_response(generator: Stream[ChatCompletionChunk]): - r"""Handle the streaming response.""" - for completion in generator: - log.debug(f"Raw chunk completion: {completion}") - parsed_content = parse_stream_response(completion) - yield parsed_content - - -def get_all_messages_content(completion: ChatCompletion) -> List[str]: - r"""When the n > 1, get all the messages content.""" - return [c.message.content for c in completion.choices] - - -def get_probabilities(completion: ChatCompletion) -> List[List[TokenLogProb]]: - r"""Get the probabilities of each token in the completion.""" - log_probs = [] - for c in completion.choices: - content = c.logprobs.content - print(content) - log_probs_for_choice = [] - for openai_token_logprob in content: - token = openai_token_logprob.token - logprob = openai_token_logprob.logprob - log_probs_for_choice.append(TokenLogProb(token=token, logprob=logprob)) - log_probs.append(log_probs_for_choice) - return log_probs - - -class OpenAIClient(ModelClient): - __doc__ = r"""A component wrapper for the OpenAI API client. - - Supports both embedding and chat completion APIs, including multimodal capabilities. - - Users can: - 1. Simplify use of ``Embedder`` and ``Generator`` components by passing `OpenAIClient()` as the `model_client`. - 2. Use this as a reference to create their own API client or extend this class by copying and modifying the code. - - Note: - We recommend avoiding `response_format` to enforce output data type or `tools` and `tool_choice` in `model_kwargs` when calling the API. - OpenAI's internal formatting and added prompts are unknown. Instead: - - Use :ref:`OutputParser` for response parsing and formatting. - - For multimodal inputs, provide images in `model_kwargs["images"]` as a path, URL, or list of them. - The model must support vision capabilities (e.g., `gpt-4o`, `gpt-4o-mini`, `o1`, `o1-mini`). - - For image generation, use `model_type=ModelType.IMAGE_GENERATION` and provide: - - model: `"dall-e-3"` or `"dall-e-2"` - - prompt: Text description of the image to generate - - size: `"1024x1024"`, `"1024x1792"`, or `"1792x1024"` for DALL-E 3; `"256x256"`, `"512x512"`, or `"1024x1024"` for DALL-E 2 - - quality: `"standard"` or `"hd"` (DALL-E 3 only) - - n: Number of images to generate (1 for DALL-E 3, 1-10 for DALL-E 2) - - response_format: `"url"` or `"b64_json"` - - Args: - api_key (Optional[str], optional): OpenAI API key. Defaults to `None`. - chat_completion_parser (Callable[[Completion], Any], optional): A function to parse the chat completion into a `str`. Defaults to `None`. - The default parser is `get_first_message_content`. - base_url (str): The API base URL to use when initializing the client. - Defaults to `"https://api.openai.com"`, but can be customized for third-party API providers or self-hosted models. - env_api_key_name (str): The environment variable name for the API key. Defaults to `"OPENAI_API_KEY"`. - - References: - - OpenAI API Overview: https://platform.openai.com/docs/introduction - - Embeddings Guide: https://platform.openai.com/docs/guides/embeddings - - Chat Completion Models: https://platform.openai.com/docs/guides/text-generation - - Vision Models: https://platform.openai.com/docs/guides/vision - - Image Generation: https://platform.openai.com/docs/guides/images - """ - - def __init__( - self, - api_key: Optional[str] = None, - chat_completion_parser: Callable[[Completion], Any] = None, - input_type: Literal["text", "messages"] = "text", - base_url: Optional[str] = None, - env_base_url_name: str = "OPENAI_BASE_URL", - env_api_key_name: str = "OPENAI_API_KEY", - ): - r"""It is recommended to set the OPENAI_API_KEY environment variable instead of passing it as an argument. - - Args: - api_key (Optional[str], optional): OpenAI API key. Defaults to None. - base_url (str): The API base URL to use when initializing the client. - env_api_key_name (str): The environment variable name for the API key. Defaults to `"OPENAI_API_KEY"`. - """ - super().__init__() - self._api_key = api_key - self._env_api_key_name = env_api_key_name - self._env_base_url_name = env_base_url_name - self.base_url = base_url or os.getenv(self._env_base_url_name, "https://api.openai.com/v1") - self.sync_client = self.init_sync_client() - self.async_client = None # only initialize if the async call is called - self.chat_completion_parser = ( - chat_completion_parser or get_first_message_content - ) - self._input_type = input_type - self._api_kwargs = {} # add api kwargs when the OpenAI Client is called - - def init_sync_client(self): - api_key = self._api_key or os.getenv(self._env_api_key_name) - if not api_key: - raise ValueError( - f"Environment variable {self._env_api_key_name} must be set" - ) - return OpenAI(api_key=api_key, base_url=self.base_url) - - def init_async_client(self): - api_key = self._api_key or os.getenv(self._env_api_key_name) - if not api_key: - raise ValueError( - f"Environment variable {self._env_api_key_name} must be set" - ) - return AsyncOpenAI(api_key=api_key, base_url=self.base_url) - - # def _parse_chat_completion(self, completion: ChatCompletion) -> "GeneratorOutput": - # # TODO: raw output it is better to save the whole completion as a source of truth instead of just the message - # try: - # data = self.chat_completion_parser(completion) - # usage = self.track_completion_usage(completion) - # return GeneratorOutput( - # data=data, error=None, raw_response=str(data), usage=usage - # ) - # except Exception as e: - # log.error(f"Error parsing the completion: {e}") - # return GeneratorOutput(data=None, error=str(e), raw_response=completion) - - def parse_chat_completion( - self, - completion: Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]], - ) -> "GeneratorOutput": - """Parse the completion, and put it into the raw_response.""" - log.debug(f"completion: {completion}, parser: {self.chat_completion_parser}") - try: - data = self.chat_completion_parser(completion) - except Exception as e: - log.error(f"Error parsing the completion: {e}") - return GeneratorOutput(data=None, error=str(e), raw_response=completion) - - try: - usage = self.track_completion_usage(completion) - return GeneratorOutput( - data=None, error=None, raw_response=data, usage=usage - ) - except Exception as e: - log.error(f"Error tracking the completion usage: {e}") - return GeneratorOutput(data=None, error=str(e), raw_response=data) - - def track_completion_usage( - self, - completion: Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]], - ) -> CompletionUsage: - - try: - usage: CompletionUsage = CompletionUsage( - completion_tokens=completion.usage.completion_tokens, - prompt_tokens=completion.usage.prompt_tokens, - total_tokens=completion.usage.total_tokens, - ) - return usage - except Exception as e: - log.error(f"Error tracking the completion usage: {e}") - return CompletionUsage( - completion_tokens=None, prompt_tokens=None, total_tokens=None - ) - - def parse_embedding_response( - self, response: CreateEmbeddingResponse - ) -> EmbedderOutput: - r"""Parse the embedding response to a structure Adalflow components can understand. - - Should be called in ``Embedder``. - """ - try: - return parse_embedding_response(response) - except Exception as e: - log.error(f"Error parsing the embedding response: {e}") - return EmbedderOutput(data=[], error=str(e), raw_response=response) - - def convert_inputs_to_api_kwargs( - self, - input: Optional[Any] = None, - model_kwargs: Dict = {}, - model_type: ModelType = ModelType.UNDEFINED, - ) -> Dict: - r""" - Specify the API input type and output api_kwargs that will be used in _call and _acall methods. - Convert the Component's standard input, and system_input(chat model) and model_kwargs into API-specific format. - For multimodal inputs, images can be provided in model_kwargs["images"] as a string path, URL, or list of them. - The model specified in model_kwargs["model"] must support multimodal capabilities when using images. - - Args: - input: The input text or messages to process - model_kwargs: Additional parameters including: - - images: Optional image source(s) as path, URL, or list of them - - detail: Image detail level ('auto', 'low', or 'high'), defaults to 'auto' - - model: The model to use (must support multimodal inputs if images are provided) - model_type: The type of model (EMBEDDER or LLM) - - Returns: - Dict: API-specific kwargs for the model call - """ - - final_model_kwargs = model_kwargs.copy() - if model_type == ModelType.EMBEDDER: - if isinstance(input, str): - input = [input] - # convert input to input - if not isinstance(input, Sequence): - raise TypeError("input must be a sequence of text") - final_model_kwargs["input"] = input - elif model_type == ModelType.LLM: - # convert input to messages - messages: List[Dict[str, str]] = [] - images = final_model_kwargs.pop("images", None) - detail = final_model_kwargs.pop("detail", "auto") - - if self._input_type == "messages": - system_start_tag = "" - system_end_tag = "" - user_start_tag = "" - user_end_tag = "" - - # new regex pattern to ignore special characters such as \n - pattern = ( - rf"{system_start_tag}\s*(.*?)\s*{system_end_tag}\s*" - rf"{user_start_tag}\s*(.*?)\s*{user_end_tag}" - ) - - # Compile the regular expression - - # re.DOTALL is to allow . to match newline so that (.*?) does not match in a single line - regex = re.compile(pattern, re.DOTALL) - # Match the pattern - match = regex.match(input) - system_prompt, input_str = None, None - - if match: - system_prompt = match.group(1) - input_str = match.group(2) - else: - print("No match found.") - if system_prompt and input_str: - messages.append({"role": "system", "content": system_prompt}) - if images: - content = [{"type": "text", "text": input_str}] - if isinstance(images, (str, dict)): - images = [images] - for img in images: - content.append(self._prepare_image_content(img, detail)) - messages.append({"role": "user", "content": content}) - else: - messages.append({"role": "user", "content": input_str}) - if len(messages) == 0: - if images: - content = [{"type": "text", "text": input}] - if isinstance(images, (str, dict)): - images = [images] - for img in images: - content.append(self._prepare_image_content(img, detail)) - messages.append({"role": "user", "content": content}) - else: - messages.append({"role": "user", "content": input}) - final_model_kwargs["messages"] = messages - elif model_type == ModelType.IMAGE_GENERATION: - # For image generation, input is the prompt - final_model_kwargs["prompt"] = input - # Ensure model is specified - if "model" not in final_model_kwargs: - raise ValueError("model must be specified for image generation") - # Set defaults for DALL-E 3 if not specified - final_model_kwargs["size"] = final_model_kwargs.get("size", "1024x1024") - final_model_kwargs["quality"] = final_model_kwargs.get( - "quality", "standard" - ) - final_model_kwargs["n"] = final_model_kwargs.get("n", 1) - final_model_kwargs["response_format"] = final_model_kwargs.get( - "response_format", "url" - ) - - # Handle image edits and variations - image = final_model_kwargs.get("image") - if isinstance(image, str) and os.path.isfile(image): - final_model_kwargs["image"] = self._encode_image(image) - - mask = final_model_kwargs.get("mask") - if isinstance(mask, str) and os.path.isfile(mask): - final_model_kwargs["mask"] = self._encode_image(mask) - else: - raise ValueError(f"model_type {model_type} is not supported") - - return final_model_kwargs - - def parse_image_generation_response(self, response: List[Image]) -> GeneratorOutput: - """Parse the image generation response into a GeneratorOutput.""" - try: - # Extract URLs or base64 data from the response - data = [img.url or img.b64_json for img in response] - # For single image responses, unwrap from list - if len(data) == 1: - data = data[0] - return GeneratorOutput( - data=data, - raw_response=str(response), - ) - except Exception as e: - log.error(f"Error parsing image generation response: {e}") - return GeneratorOutput(data=None, error=str(e), raw_response=str(response)) - - @backoff.on_exception( - backoff.expo, - ( - APITimeoutError, - InternalServerError, - RateLimitError, - UnprocessableEntityError, - BadRequestError, - ), - max_time=5, - ) - def call(self, api_kwargs: Dict = {}, model_type: ModelType = ModelType.UNDEFINED): - """ - kwargs is the combined input and model_kwargs. Support streaming call. - """ - log.info(f"api_kwargs: {api_kwargs}") - self._api_kwargs = api_kwargs - if model_type == ModelType.EMBEDDER: - return self.sync_client.embeddings.create(**api_kwargs) - elif model_type == ModelType.LLM: - if "stream" in api_kwargs and api_kwargs.get("stream", False): - log.debug("streaming call") - self.chat_completion_parser = handle_streaming_response - return self.sync_client.chat.completions.create(**api_kwargs) - else: - log.debug("non-streaming call converted to streaming") - # Make a copy of api_kwargs to avoid modifying the original - streaming_kwargs = api_kwargs.copy() - streaming_kwargs["stream"] = True - - # Get streaming response - stream_response = self.sync_client.chat.completions.create(**streaming_kwargs) - - # Accumulate all content from the stream - accumulated_content = "" - id = "" - model = "" - created = 0 - for chunk in stream_response: - id = getattr(chunk, "id", None) or id - model = getattr(chunk, "model", None) or model - created = getattr(chunk, "created", 0) or created - choices = getattr(chunk, "choices", []) - if len(choices) > 0: - delta = getattr(choices[0], "delta", None) - if delta is not None: - text = getattr(delta, "content", None) - if text is not None: - accumulated_content += text or "" - # Return the mock completion object that will be processed by the chat_completion_parser - return ChatCompletion( - id = id, - model=model, - created=created, - object="chat.completion", - choices=[Choice( - index=0, - finish_reason="stop", - message=ChatCompletionMessage(content=accumulated_content, role="assistant") - )] - ) - elif model_type == ModelType.IMAGE_GENERATION: - # Determine which image API to call based on the presence of image/mask - if "image" in api_kwargs: - if "mask" in api_kwargs: - # Image edit - response = self.sync_client.images.edit(**api_kwargs) - else: - # Image variation - response = self.sync_client.images.create_variation(**api_kwargs) - else: - # Image generation - response = self.sync_client.images.generate(**api_kwargs) - return response.data - else: - raise ValueError(f"model_type {model_type} is not supported") - - @backoff.on_exception( - backoff.expo, - ( - APITimeoutError, - InternalServerError, - RateLimitError, - UnprocessableEntityError, - BadRequestError, - ), - max_time=5, - ) - async def acall( - self, api_kwargs: Dict = {}, model_type: ModelType = ModelType.UNDEFINED - ): - """ - kwargs is the combined input and model_kwargs - """ - # store the api kwargs in the client - self._api_kwargs = api_kwargs - if self.async_client is None: - self.async_client = self.init_async_client() - if model_type == ModelType.EMBEDDER: - return await self.async_client.embeddings.create(**api_kwargs) - elif model_type == ModelType.LLM: - return await self.async_client.chat.completions.create(**api_kwargs) - elif model_type == ModelType.IMAGE_GENERATION: - # Determine which image API to call based on the presence of image/mask - if "image" in api_kwargs: - if "mask" in api_kwargs: - # Image edit - response = await self.async_client.images.edit(**api_kwargs) - else: - # Image variation - response = await self.async_client.images.create_variation( - **api_kwargs - ) - else: - # Image generation - response = await self.async_client.images.generate(**api_kwargs) - return response.data - else: - raise ValueError(f"model_type {model_type} is not supported") - - @classmethod - def from_dict(cls: type[T], data: Dict[str, Any]) -> T: - obj = super().from_dict(data) - # recreate the existing clients - obj.sync_client = obj.init_sync_client() - obj.async_client = obj.init_async_client() - return obj - - def to_dict(self) -> Dict[str, Any]: - r"""Convert the component to a dictionary.""" - # TODO: not exclude but save yes or no for recreating the clients - exclude = [ - "sync_client", - "async_client", - ] # unserializable object - output = super().to_dict(exclude=exclude) - return output - - def _encode_image(self, image_path: str) -> str: - """Encode image to base64 string. - - Args: - image_path: Path to image file. - - Returns: - Base64 encoded image string. - - Raises: - ValueError: If the file cannot be read or doesn't exist. - """ - try: - with open(image_path, "rb") as image_file: - return base64.b64encode(image_file.read()).decode("utf-8") - except FileNotFoundError: - raise ValueError(f"Image file not found: {image_path}") - except PermissionError: - raise ValueError(f"Permission denied when reading image file: {image_path}") - except Exception as e: - raise ValueError(f"Error encoding image {image_path}: {str(e)}") - - def _prepare_image_content( - self, image_source: Union[str, Dict[str, Any]], detail: str = "auto" - ) -> Dict[str, Any]: - """Prepare image content for API request. - - Args: - image_source: Either a path to local image or a URL. - detail: Image detail level ('auto', 'low', or 'high'). - - Returns: - Formatted image content for API request. - """ - if isinstance(image_source, str): - if image_source.startswith(("http://", "https://")): - return { - "type": "image_url", - "image_url": {"url": image_source, "detail": detail}, - } - else: - base64_image = self._encode_image(image_source) - return { - "type": "image_url", - "image_url": { - "url": f"data:image/jpeg;base64,{base64_image}", - "detail": detail, - }, - } - return image_source - - -# Example usage: -if __name__ == "__main__": - from adalflow.core import Generator - from adalflow.utils import setup_env - - # log = get_logger(level="DEBUG") - - setup_env() - prompt_kwargs = {"input_str": "What is the meaning of life?"} - - gen = Generator( - model_client=OpenAIClient(), - model_kwargs={"model": "gpt-4o", "stream": False}, - ) - gen_response = gen(prompt_kwargs) - print(f"gen_response: {gen_response}") - - # for genout in gen_response.data: - # print(f"genout: {genout}") - - # test that to_dict and from_dict works - # model_client = OpenAIClient() - # model_client_dict = model_client.to_dict() - # from_dict_model_client = OpenAIClient.from_dict(model_client_dict) - # assert model_client_dict == from_dict_model_client.to_dict() - - -if __name__ == "__main__": - import adalflow as adal - - # setup env or pass the api_key - from adalflow.utils import setup_env - - setup_env() - - openai_llm = adal.Generator( - model_client=OpenAIClient(), model_kwargs={"model": "gpt-4o"} - ) - resopnse = openai_llm(prompt_kwargs={"input_str": "What is LLM?"}) - print(resopnse) From c411bee046038661e8b3cc8f82978702d8afeb96 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 19 Jul 2026 10:45:30 +0800 Subject: [PATCH 09/25] centralize prompt (related to #306) --- api/websocket_wiki.py | 163 +++++++++--------------------------------- 1 file changed, 32 insertions(+), 131 deletions(-) diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index 7064f54b4..ebf23486f 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -26,6 +26,12 @@ from api.azureai_client import AzureAIClient from api.dashscope_client import DashscopeClient from api.rag import RAG +from api.prompts import ( + DEEP_RESEARCH_FIRST_ITERATION_PROMPT, + DEEP_RESEARCH_FINAL_ITERATION_PROMPT, + DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, + SIMPLE_CHAT_SYSTEM_PROMPT, +) # Configure logging from api.logging_config import setup_logging @@ -265,140 +271,35 @@ async def handle_websocket_chat(websocket: WebSocket): is_final_iteration = research_iteration >= 5 if is_first_iteration: - system_prompt = f""" -You are an expert code analyst examining the {repo_type} repository: {repo_url} ({repo_name}). -You are conducting a multi-turn Deep Research process to thoroughly investigate the specific topic in the user's query. -Your goal is to provide detailed, focused information EXCLUSIVELY about this topic. -IMPORTANT:You MUST respond in {language_name} language. - - - -- This is the first iteration of a multi-turn research process focused EXCLUSIVELY on the user's query -- Start your response with "## Research Plan" -- Outline your approach to investigating this specific topic -- If the topic is about a specific file or feature (like "Dockerfile"), focus ONLY on that file or feature -- Clearly state the specific topic you're researching to maintain focus throughout all iterations -- Identify the key aspects you'll need to research -- Provide initial findings based on the information available -- End with "## Next Steps" indicating what you'll investigate in the next iteration -- Do NOT provide a final conclusion yet - this is just the beginning of the research -- Do NOT include general repository information unless directly relevant to the query -- Focus EXCLUSIVELY on the specific topic being researched - do not drift to related topics -- Your research MUST directly address the original question -- NEVER respond with just "Continue the research" as an answer - always provide substantive research findings -- Remember that this topic will be maintained across all research iterations - - -""" + system_prompt = DEEP_RESEARCH_FIRST_ITERATION_PROMPT.format( + repo_type=repo_type, + repo_url=repo_url, + repo_name=repo_name, + language_name=language_name + ) elif is_final_iteration: - system_prompt = f""" -You are an expert code analyst examining the {repo_type} repository: {repo_url} ({repo_name}). -You are in the final iteration of a Deep Research process focused EXCLUSIVELY on the latest user query. -Your goal is to synthesize all previous findings and provide a comprehensive conclusion that directly addresses this specific topic and ONLY this topic. -IMPORTANT:You MUST respond in {language_name} language. - - - -- This is the final iteration of the research process -- CAREFULLY review the entire conversation history to understand all previous findings -- Synthesize ALL findings from previous iterations into a comprehensive conclusion -- Start with "## Final Conclusion" -- Your conclusion MUST directly address the original question -- Stay STRICTLY focused on the specific topic - do not drift to related topics -- Include specific code references and implementation details related to the topic -- Highlight the most important discoveries and insights about this specific functionality -- Provide a complete and definitive answer to the original question -- Do NOT include general repository information unless directly relevant to the query -- Focus exclusively on the specific topic being researched -- NEVER respond with "Continue the research" as an answer - always provide a complete conclusion -- If the topic is about a specific file or feature (like "Dockerfile"), focus ONLY on that file or feature -- Ensure your conclusion builds on and references key findings from previous iterations - - -""" + system_prompt = DEEP_RESEARCH_FINAL_ITERATION_PROMPT.format( + repo_type=repo_type, + repo_url=repo_url, + repo_name=repo_name, + research_iteration=research_iteration, + language_name=language_name + ) else: - system_prompt = f""" -You are an expert code analyst examining the {repo_type} repository: {repo_url} ({repo_name}). -You are currently in iteration {research_iteration} of a Deep Research process focused EXCLUSIVELY on the latest user query. -Your goal is to build upon previous research iterations and go deeper into this specific topic without deviating from it. -IMPORTANT:You MUST respond in {language_name} language. - - - -- CAREFULLY review the conversation history to understand what has been researched so far -- Your response MUST build on previous research iterations - do not repeat information already covered -- Identify gaps or areas that need further exploration related to this specific topic -- Focus on one specific aspect that needs deeper investigation in this iteration -- Start your response with "## Research Update {research_iteration}" -- Clearly explain what you're investigating in this iteration -- Provide new insights that weren't covered in previous iterations -- If this is iteration 3, prepare for a final conclusion in the next iteration -- Do NOT include general repository information unless directly relevant to the query -- Focus EXCLUSIVELY on the specific topic being researched - do not drift to related topics -- If the topic is about a specific file or feature (like "Dockerfile"), focus ONLY on that file or feature -- NEVER respond with just "Continue the research" as an answer - always provide substantive research findings -- Your research MUST directly address the original question -- Maintain continuity with previous research iterations - this is a continuous investigation - - -""" + system_prompt = DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT.format( + repo_type=repo_type, + repo_url=repo_url, + repo_name=repo_name, + research_iteration=research_iteration, + language_name=language_name + ) else: - system_prompt = f""" -You are an expert code analyst examining the {repo_type} repository: {repo_url} ({repo_name}). -You provide direct, concise, and accurate information about code repositories. -You NEVER start responses with markdown headers or code fences. -IMPORTANT:You MUST respond in {language_name} language. - - - -- Answer the user's question directly without ANY preamble or filler phrases -- DO NOT include any rationale, explanation, or extra comments. -- Strictly base answers ONLY on existing code or documents -- DO NOT speculate or invent citations. -- DO NOT start with preambles like "Okay, here's a breakdown" or "Here's an explanation" -- DO NOT start with markdown headers like "## Analysis of..." or any file path references -- DO NOT start with ```markdown code fences -- DO NOT end your response with ``` closing fences -- DO NOT start by repeating or acknowledging the question -- JUST START with the direct answer to the question - - -```markdown -## Analysis of `adalflow/adalflow/datasets/gsm8k.py` - -This file contains... -``` - - -- Format your response with proper markdown including headings, lists, and code blocks WITHIN your answer -- For code analysis, organize your response with clear sections -- Think step by step and structure your answer logically -- Start with the most relevant information that directly addresses the user's query -- Be precise and technical when discussing code -- Your response language should be in the same language as the user's query - - -""" + system_prompt = SIMPLE_CHAT_SYSTEM_PROMPT.format( + repo_type=repo_type, + repo_url=repo_url, + repo_name=repo_name, + language_name=language_name + ) # Fetch file content if provided file_content = "" From 0cb89571642459ea2f73aa3429e4eee32560b757 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 19 Jul 2026 10:48:12 +0800 Subject: [PATCH 10/25] centralize prompt (related to #306) (#8) --- api/websocket_wiki.py | 163 +++++++++--------------------------------- 1 file changed, 32 insertions(+), 131 deletions(-) diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index 87d508726..ca3fb04fd 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -14,6 +14,12 @@ ) from api.data_pipeline import count_tokens, get_file_content from api.rag import RAG +from api.prompts import ( + DEEP_RESEARCH_FIRST_ITERATION_PROMPT, + DEEP_RESEARCH_FINAL_ITERATION_PROMPT, + DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, + SIMPLE_CHAT_SYSTEM_PROMPT, +) # Configure logging from api.logging_config import setup_logging @@ -253,140 +259,35 @@ async def handle_websocket_chat(websocket: WebSocket): is_final_iteration = research_iteration >= 5 if is_first_iteration: - system_prompt = f""" -You are an expert code analyst examining the {repo_type} repository: {repo_url} ({repo_name}). -You are conducting a multi-turn Deep Research process to thoroughly investigate the specific topic in the user's query. -Your goal is to provide detailed, focused information EXCLUSIVELY about this topic. -IMPORTANT:You MUST respond in {language_name} language. - - - -- This is the first iteration of a multi-turn research process focused EXCLUSIVELY on the user's query -- Start your response with "## Research Plan" -- Outline your approach to investigating this specific topic -- If the topic is about a specific file or feature (like "Dockerfile"), focus ONLY on that file or feature -- Clearly state the specific topic you're researching to maintain focus throughout all iterations -- Identify the key aspects you'll need to research -- Provide initial findings based on the information available -- End with "## Next Steps" indicating what you'll investigate in the next iteration -- Do NOT provide a final conclusion yet - this is just the beginning of the research -- Do NOT include general repository information unless directly relevant to the query -- Focus EXCLUSIVELY on the specific topic being researched - do not drift to related topics -- Your research MUST directly address the original question -- NEVER respond with just "Continue the research" as an answer - always provide substantive research findings -- Remember that this topic will be maintained across all research iterations - - -""" + system_prompt = DEEP_RESEARCH_FIRST_ITERATION_PROMPT.format( + repo_type=repo_type, + repo_url=repo_url, + repo_name=repo_name, + language_name=language_name + ) elif is_final_iteration: - system_prompt = f""" -You are an expert code analyst examining the {repo_type} repository: {repo_url} ({repo_name}). -You are in the final iteration of a Deep Research process focused EXCLUSIVELY on the latest user query. -Your goal is to synthesize all previous findings and provide a comprehensive conclusion that directly addresses this specific topic and ONLY this topic. -IMPORTANT:You MUST respond in {language_name} language. - - - -- This is the final iteration of the research process -- CAREFULLY review the entire conversation history to understand all previous findings -- Synthesize ALL findings from previous iterations into a comprehensive conclusion -- Start with "## Final Conclusion" -- Your conclusion MUST directly address the original question -- Stay STRICTLY focused on the specific topic - do not drift to related topics -- Include specific code references and implementation details related to the topic -- Highlight the most important discoveries and insights about this specific functionality -- Provide a complete and definitive answer to the original question -- Do NOT include general repository information unless directly relevant to the query -- Focus exclusively on the specific topic being researched -- NEVER respond with "Continue the research" as an answer - always provide a complete conclusion -- If the topic is about a specific file or feature (like "Dockerfile"), focus ONLY on that file or feature -- Ensure your conclusion builds on and references key findings from previous iterations - - -""" + system_prompt = DEEP_RESEARCH_FINAL_ITERATION_PROMPT.format( + repo_type=repo_type, + repo_url=repo_url, + repo_name=repo_name, + research_iteration=research_iteration, + language_name=language_name + ) else: - system_prompt = f""" -You are an expert code analyst examining the {repo_type} repository: {repo_url} ({repo_name}). -You are currently in iteration {research_iteration} of a Deep Research process focused EXCLUSIVELY on the latest user query. -Your goal is to build upon previous research iterations and go deeper into this specific topic without deviating from it. -IMPORTANT:You MUST respond in {language_name} language. - - - -- CAREFULLY review the conversation history to understand what has been researched so far -- Your response MUST build on previous research iterations - do not repeat information already covered -- Identify gaps or areas that need further exploration related to this specific topic -- Focus on one specific aspect that needs deeper investigation in this iteration -- Start your response with "## Research Update {research_iteration}" -- Clearly explain what you're investigating in this iteration -- Provide new insights that weren't covered in previous iterations -- If this is iteration 3, prepare for a final conclusion in the next iteration -- Do NOT include general repository information unless directly relevant to the query -- Focus EXCLUSIVELY on the specific topic being researched - do not drift to related topics -- If the topic is about a specific file or feature (like "Dockerfile"), focus ONLY on that file or feature -- NEVER respond with just "Continue the research" as an answer - always provide substantive research findings -- Your research MUST directly address the original question -- Maintain continuity with previous research iterations - this is a continuous investigation - - -""" + system_prompt = DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT.format( + repo_type=repo_type, + repo_url=repo_url, + repo_name=repo_name, + research_iteration=research_iteration, + language_name=language_name + ) else: - system_prompt = f""" -You are an expert code analyst examining the {repo_type} repository: {repo_url} ({repo_name}). -You provide direct, concise, and accurate information about code repositories. -You NEVER start responses with markdown headers or code fences. -IMPORTANT:You MUST respond in {language_name} language. - - - -- Answer the user's question directly without ANY preamble or filler phrases -- DO NOT include any rationale, explanation, or extra comments. -- Strictly base answers ONLY on existing code or documents -- DO NOT speculate or invent citations. -- DO NOT start with preambles like "Okay, here's a breakdown" or "Here's an explanation" -- DO NOT start with markdown headers like "## Analysis of..." or any file path references -- DO NOT start with ```markdown code fences -- DO NOT end your response with ``` closing fences -- DO NOT start by repeating or acknowledging the question -- JUST START with the direct answer to the question - - -```markdown -## Analysis of `adalflow/adalflow/datasets/gsm8k.py` - -This file contains... -``` - - -- Format your response with proper markdown including headings, lists, and code blocks WITHIN your answer -- For code analysis, organize your response with clear sections -- Think step by step and structure your answer logically -- Start with the most relevant information that directly addresses the user's query -- Be precise and technical when discussing code -- Your response language should be in the same language as the user's query - - -""" + system_prompt = SIMPLE_CHAT_SYSTEM_PROMPT.format( + repo_type=repo_type, + repo_url=repo_url, + repo_name=repo_name, + language_name=language_name + ) # Fetch file content if provided file_content = "" From 075fbd555f00d2ba8685284c046befa17a0b9e7e Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 19 Jul 2026 12:59:40 +0800 Subject: [PATCH 11/25] fix prompt --- api/prompts.py | 2 +- api/simple_chat.py | 1 - api/websocket_wiki.py | 1 - 3 files changed, 1 insertion(+), 3 deletions(-) diff --git a/api/prompts.py b/api/prompts.py index 61ef0a4d3..7c6641553 100644 --- a/api/prompts.py +++ b/api/prompts.py @@ -131,7 +131,7 @@ - Your response MUST build on previous research iterations - do not repeat information already covered - Identify gaps or areas that need further exploration related to this specific topic - Focus on one specific aspect that needs deeper investigation in this iteration -- Start your response with "## Research Update {{research_iteration}}" +- Start your response with "## Research Update {research_iteration}" - Clearly explain what you're investigating in this iteration - Provide new insights that weren't covered in previous iterations - If this is iteration 3, prepare for a final conclusion in the next iteration diff --git a/api/simple_chat.py b/api/simple_chat.py index e93199397..48204c9ba 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -273,7 +273,6 @@ async def chat_completions_stream(request: ChatCompletionRequest): repo_type=repo_type, repo_url=repo_url, repo_name=repo_name, - research_iteration=research_iteration, language_name=language_name ) else: diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index ebf23486f..67b1ba525 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -282,7 +282,6 @@ async def handle_websocket_chat(websocket: WebSocket): repo_type=repo_type, repo_url=repo_url, repo_name=repo_name, - research_iteration=research_iteration, language_name=language_name ) else: From 084358875f8b1ea002998ae9cddc68480cf92eed Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 19 Jul 2026 15:32:53 +0800 Subject: [PATCH 12/25] update `check_ollama_model_exists` using ollama.list api. --- api/ollama_patch.py | 28 +++++++++++----------------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/api/ollama_patch.py b/api/ollama_patch.py index bb6a77294..a7b4cb482 100644 --- a/api/ollama_patch.py +++ b/api/ollama_patch.py @@ -23,30 +23,24 @@ def check_ollama_model_exists(model_name: str, ollama_host: str = None) -> bool: Returns: bool: True if model exists, False otherwise """ + import ollama if ollama_host is None: ollama_host = os.getenv("OLLAMA_HOST", "http://localhost:11434") - try: # Remove /api prefix if present and add it back if ollama_host.endswith('/api'): ollama_host = ollama_host[:-4] - - response = requests.get(f"{ollama_host}/api/tags", timeout=5) - if response.status_code == 200: - models_data = response.json() - available_models = [model.get('name', '').split(':')[0] for model in models_data.get('models', [])] - model_base_name = model_name.split(':')[0] # Remove tag if present - - is_available = model_base_name in available_models - if is_available: - logger.info(f"Ollama model '{model_name}' is available") - else: - logger.warning(f"Ollama model '{model_name}' is not available. Available models: {available_models}") - return is_available + ret: ollama.ListResponse = ollama.Client(host=ollama_host, timeout=5).list() + is_available = any(model_name == model.model for model in ret.models) + if is_available: + logger.info("Ollama model '%s' is available", model_name) else: - logger.warning(f"Could not check Ollama models, status code: {response.status_code}") - return False - except requests.exceptions.RequestException as e: + logger.warning( + "Ollama model '%s' is not available. Available models: %s. ", + model_name, + str([model.model for model in ret.models])) + return is_available + except ConnectionError as e: logger.warning(f"Could not connect to Ollama to check models: {e}") return False except Exception as e: From da715858bca072b964fe26d072432d46ac15de4b Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 19 Jul 2026 15:45:08 +0800 Subject: [PATCH 13/25] fix exception type --- api/ollama_patch.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/api/ollama_patch.py b/api/ollama_patch.py index a7b4cb482..9ba5ac34f 100644 --- a/api/ollama_patch.py +++ b/api/ollama_patch.py @@ -1,4 +1,6 @@ import logging + +import httpx import requests import os @@ -40,7 +42,7 @@ def check_ollama_model_exists(model_name: str, ollama_host: str = None) -> bool: model_name, str([model.model for model in ret.models])) return is_available - except ConnectionError as e: + except (httpx.ConnectTimeout, ConnectionError): logger.warning(f"Could not connect to Ollama to check models: {e}") return False except Exception as e: From 592523e581de456fb9b725b1d6a1effbc50458fe Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 19 Jul 2026 15:46:05 +0800 Subject: [PATCH 14/25] fix import scope --- api/ollama_patch.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/api/ollama_patch.py b/api/ollama_patch.py index 9ba5ac34f..ab07b033f 100644 --- a/api/ollama_patch.py +++ b/api/ollama_patch.py @@ -1,7 +1,5 @@ import logging -import httpx -import requests import os # Configure logging @@ -26,6 +24,7 @@ def check_ollama_model_exists(model_name: str, ollama_host: str = None) -> bool: bool: True if model exists, False otherwise """ import ollama + import httpx if ollama_host is None: ollama_host = os.getenv("OLLAMA_HOST", "http://localhost:11434") try: From 84f051723edbb8d6aab2060a484fa4d3213b3a39 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 19 Jul 2026 15:51:31 +0800 Subject: [PATCH 15/25] fix catching exception --- api/ollama_patch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/ollama_patch.py b/api/ollama_patch.py index ab07b033f..cc04b7317 100644 --- a/api/ollama_patch.py +++ b/api/ollama_patch.py @@ -41,7 +41,7 @@ def check_ollama_model_exists(model_name: str, ollama_host: str = None) -> bool: model_name, str([model.model for model in ret.models])) return is_available - except (httpx.ConnectTimeout, ConnectionError): + except (httpx.ConnectTimeout, ConnectionError) as e: logger.warning(f"Could not connect to Ollama to check models: {e}") return False except Exception as e: From 0bf2d4b079a30410bab9d9cc754f0cdac91fc8e7 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sun, 19 Jul 2026 19:52:45 +0800 Subject: [PATCH 16/25] remove dangling file. --- api/ollama_client.py | 131 ------------------------------------------- 1 file changed, 131 deletions(-) delete mode 100644 api/ollama_client.py diff --git a/api/ollama_client.py b/api/ollama_client.py deleted file mode 100644 index 2418e313d..000000000 --- a/api/ollama_client.py +++ /dev/null @@ -1,131 +0,0 @@ -# This file exists to patch the adalflow OllamaClient to support ollama batch embedding api `embed` - -from typing import Dict, Optional, Any - -from adalflow.components.model_client.ollama_client import ( - OllamaClient, - RequestError, - ResponseError, - log, -) -from adalflow.core.types import EmbedderOutput, Embedding - -from adalflow.core import ModelType - -import backoff - - -def convert_inputs_to_api_kwargs( - self, - input: Optional[Any] = None, - model_kwargs: Dict = {}, - model_type: ModelType = ModelType.UNDEFINED, -) -> Dict: - self.generate = False - final_model_kwargs = model_kwargs.copy() - if model_type == ModelType.EMBEDDER: - final_model_kwargs["input"] = input - return final_model_kwargs - elif model_type == ModelType.LLM: - if input is not None and input != "": - # check if "generate" is in model_kwargs, and if it is set to True, then we use generate api - if model_kwargs.get("generate", False): - final_model_kwargs["prompt"] = input - self.generate = True - else: - # for chat api, we need to convert the input to a message format - # if the input is a string, we create a message with role "user" - if isinstance(input, str): - input = [{"role": "user", "content": input}] - final_model_kwargs["messages"] = input - elif not isinstance(input, list): - raise ValueError("Input must be a string or a list of messages") - # if the input is a list of messages, we use it as is - return final_model_kwargs - else: - raise ValueError("Input must be text") - else: - raise ValueError(f"model_type {model_type} is not supported") - - -@backoff.on_exception( - backoff.expo, - (RequestError, ResponseError), - max_tries=5, -) -async def acall( - self, - api_kwargs: dict = None, - model_type: ModelType = ModelType.UNDEFINED, -): - if self.async_client is None: - self.init_async_client() - if self.async_client is None: - raise RuntimeError( - "Async client is not initialized" - ) - api_kwargs = api_kwargs or {} - if model_type == ModelType.EMBEDDER: - return await self.async_client.embed(**api_kwargs) - if model_type == ModelType.LLM: # in default we use chat - # create a message from the input - if "generate" in api_kwargs and api_kwargs["generate"]: - # remove generate from api_kwargs - api_kwargs.pop("generate") - return await self.async_client.generate(**api_kwargs) - else: - return await self.async_client.chat(**api_kwargs) - else: - raise ValueError(f"model_type {model_type} is not supported") - - - -@backoff.on_exception( - backoff.expo, - (RequestError, ResponseError), - max_tries=5, -) -def call( - self, - api_kwargs: dict = None, - model_type: ModelType = ModelType.UNDEFINED, -): - api_kwargs = api_kwargs or {} - if not self.sync_client: - self.init_sync_client() - if self.sync_client is None: - raise RuntimeError("Sync client is not initialized") - - if model_type == ModelType.EMBEDDER: - return self.sync_client.embed(**api_kwargs) - if model_type == ModelType.LLM: - if "generate" in api_kwargs and api_kwargs["generate"]: - # remove generate from api_kwargs - api_kwargs.pop("generate") - return self.sync_client.generate(**api_kwargs) - else: - return self.sync_client.chat(**api_kwargs) - else: - raise ValueError(f"model_type {model_type} is not supported") - - -def parse_embedding_response( - self, response: Dict[str, list[float]] -) -> EmbedderOutput: - r"""Parse the embedding response to a structure AdalFlow components can understand. - Pull the embedding from response['embedding'] and store it Embedding dataclass - """ - try: - return EmbedderOutput(data=[ - Embedding(embedding=emb, index=i) - for i, emb in enumerate(response["embeddings"]) - ]) - except Exception as e: - log.error(f"Error parsing the embedding response: {e}") - return EmbedderOutput(data=[], error=str(e), raw_response=response) - -OllamaClient.convert_inputs_to_api_kwargs = convert_inputs_to_api_kwargs -OllamaClient.parse_embedding_response = parse_embedding_response -OllamaClient.call = call -OllamaClient.acall = acall - From 9907fdc1db0f9eee9dd74418eca990aab3296411 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 09:40:02 +0800 Subject: [PATCH 17/25] use field_validator to simplify dirs and files parsing in `ChatCompletionRequest` --- api/simple_chat.py | 71 ++++++++++++++++++++++++++++++++-------------- 1 file changed, 49 insertions(+), 22 deletions(-) diff --git a/api/simple_chat.py b/api/simple_chat.py index 9076f71ab..4b48ff3fc 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -1,12 +1,12 @@ import logging -from typing import List, Optional, Callable +from typing import List, Optional, Callable, Literal from urllib.parse import unquote from functools import partial from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator from api.chat import ChatStreamer, prompt_builder, is_token_limit_error from api.config import get_model_config, configs @@ -54,17 +54,45 @@ class ChatCompletionRequest(BaseModel): messages: List[ChatMessage] = Field(..., description="List of chat messages") filePath: Optional[str] = Field(None, description="Optional path to a file in the repository to include in the prompt") token: Optional[str] = Field(None, description="Personal access token for private repositories") - type: Optional[str] = Field("github", description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')") + type: Optional[Literal["github", "gitlab", "bitbucket"]] = Field( + "github", + description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')", + ) # model parameters provider: str = Field("google", description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)") model: Optional[str] = Field(None, description="Model name for the specified provider") language: Optional[str] = Field("en", description="Language for content generation (e.g., 'en', 'ja', 'zh', 'es', 'kr', 'vi')") - excluded_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to exclude from processing") - excluded_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to exclude from processing") - included_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to include exclusively") - included_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to include exclusively") + excluded_dirs: List[str] = Field( + default_factory=list, + description="Comma-separated list of directories to exclude from processing", + ) + excluded_files: List[str] = Field( + default_factory=list, + description="Comma-separated list of file patterns to exclude from processing", + ) + included_dirs: List[str] = Field( + default_factory=list, + description="Comma-separated list of directories to include exclusively", + ) + included_files: List[str] = Field( + default_factory=list, + description="Comma-separated list of file patterns to include exclusively", + ) + + @field_validator( + "excluded_dirs", + "excluded_files", + "included_dirs", + "included_files", + mode="before", + ) + @classmethod + def validate_path(cls, value: list[str] | str) -> list[str]: + if isinstance(value, str): + value = [unquote(path) for path in value.strip().split("\n")] + return value @app.post("/chat/completions/stream") @@ -87,25 +115,24 @@ async def chat_completions_stream(request: ChatCompletionRequest): request_rag = RAG(provider=request.provider, model=request.model) # Extract custom file filter parameters if provided - excluded_dirs = None - excluded_files = None - included_dirs = None - included_files = None - if request.excluded_dirs: - excluded_dirs = [unquote(dir_path) for dir_path in request.excluded_dirs.split('\n') if dir_path.strip()] - logger.info(f"Using custom excluded directories: {excluded_dirs}") + logger.info(f"Using custom excluded directories: {request.excluded_dirs}") if request.excluded_files: - excluded_files = [unquote(file_pattern) for file_pattern in request.excluded_files.split('\n') if file_pattern.strip()] - logger.info(f"Using custom excluded files: {excluded_files}") + logger.info(f"Using custom excluded files: {request.excluded_files}") if request.included_dirs: - included_dirs = [unquote(dir_path) for dir_path in request.included_dirs.split('\n') if dir_path.strip()] - logger.info(f"Using custom included directories: {included_dirs}") + logger.info(f"Using custom included directories: {request.included_dirs}") if request.included_files: - included_files = [unquote(file_pattern) for file_pattern in request.included_files.split('\n') if file_pattern.strip()] - logger.info(f"Using custom included files: {included_files}") - - request_rag.prepare_retriever(request.repo_url, request.type, request.token, excluded_dirs, excluded_files, included_dirs, included_files) + logger.info(f"Using custom included files: {request.included_files}") + + request_rag.prepare_retriever( + request.repo_url, + request.type, + request.token, + excluded_dirs=request.excluded_dirs, + excluded_files=request.excluded_files, + included_dirs=request.included_dirs, + included_files=request.included_files, + ) logger.info(f"Retriever prepared for {request.repo_url}") except ValueError as e: if "No valid documents with embeddings found" in str(e): From e66ae526c929e116ec29437f10d5e723797f74fd Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 09:46:20 +0800 Subject: [PATCH 18/25] Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- api/simple_chat.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/api/simple_chat.py b/api/simple_chat.py index 4b48ff3fc..e37433643 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -66,19 +66,19 @@ class ChatCompletionRequest(BaseModel): language: Optional[str] = Field("en", description="Language for content generation (e.g., 'en', 'ja', 'zh', 'es', 'kr', 'vi')") excluded_dirs: List[str] = Field( default_factory=list, - description="Comma-separated list of directories to exclude from processing", + description="List or newline-separated string of directories to exclude from processing", ) excluded_files: List[str] = Field( default_factory=list, - description="Comma-separated list of file patterns to exclude from processing", + description="List or newline-separated string of file patterns to exclude from processing", ) included_dirs: List[str] = Field( default_factory=list, - description="Comma-separated list of directories to include exclusively", + description="List or newline-separated string of directories to include exclusively", ) included_files: List[str] = Field( default_factory=list, - description="Comma-separated list of file patterns to include exclusively", + description="List or newline-separated string of file patterns to include exclusively", ) @field_validator( From 990f7b9fecec4d77c77af12ea966d351ed37a453 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 09:51:25 +0800 Subject: [PATCH 19/25] prevent empty string --- api/simple_chat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/simple_chat.py b/api/simple_chat.py index e37433643..23a9cd921 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -91,7 +91,7 @@ class ChatCompletionRequest(BaseModel): @classmethod def validate_path(cls, value: list[str] | str) -> list[str]: if isinstance(value, str): - value = [unquote(path) for path in value.strip().split("\n")] + value = [unquote(path) for path in value.strip().split("\n") if path] return value From c233a4420d6e35314491dd613ebfb35e82f7e4b5 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 10:26:26 +0800 Subject: [PATCH 20/25] unify chat model in `websocket_wiki` and `simple_chat` --- api/chat_model.py | 57 ++++++++++++++++++++++++++++++++++++++++ api/simple_chat.py | 55 +------------------------------------- api/websocket_wiki.py | 61 +++++++++++-------------------------------- 3 files changed, 73 insertions(+), 100 deletions(-) create mode 100644 api/chat_model.py diff --git a/api/chat_model.py b/api/chat_model.py new file mode 100644 index 000000000..f7cbf2fc4 --- /dev/null +++ b/api/chat_model.py @@ -0,0 +1,57 @@ +from typing import Optional, List, Literal +from urllib.parse import unquote + +from pydantic import BaseModel, Field, field_validator + +# Models for the API +class ChatMessage(BaseModel): + role: str # 'user' or 'assistant' + content: str + +class ChatCompletionRequest(BaseModel): + """ + Model for requesting a chat completion. + """ + repo_url: str = Field(..., description="URL of the repository to query") + messages: List[ChatMessage] = Field(..., description="List of chat messages") + filePath: Optional[str] = Field(None, description="Optional path to a file in the repository to include in the prompt") + token: Optional[str] = Field(None, description="Personal access token for private repositories") + type: Optional[Literal["github", "gitlab", "bitbucket"]] = Field( + "github", + description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')", + ) + + # model parameters + provider: str = Field("google", description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)") + model: Optional[str] = Field(None, description="Model name for the specified provider") + + language: Optional[str] = Field("en", description="Language for content generation (e.g., 'en', 'ja', 'zh', 'es', 'kr', 'vi')") + excluded_dirs: List[str] = Field( + default_factory=list, + description="List or newline-separated string of directories to exclude from processing", + ) + excluded_files: List[str] = Field( + default_factory=list, + description="List or newline-separated string of file patterns to exclude from processing", + ) + included_dirs: List[str] = Field( + default_factory=list, + description="List or newline-separated string of directories to include exclusively", + ) + included_files: List[str] = Field( + default_factory=list, + description="List or newline-separated string of file patterns to include exclusively", + ) + + @field_validator( + "excluded_dirs", + "excluded_files", + "included_dirs", + "included_files", + mode="before", + ) + @classmethod + def validate_path(cls, value: list[str] | str) -> list[str]: + if isinstance(value, str): + value = [unquote(path) for path in value.strip().split("\n") if path] + return value diff --git a/api/simple_chat.py b/api/simple_chat.py index 23a9cd921..dec70345a 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -6,7 +6,6 @@ from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse -from pydantic import BaseModel, Field, field_validator from api.chat import ChatStreamer, prompt_builder, is_token_limit_error from api.config import get_model_config, configs @@ -18,6 +17,7 @@ DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, SIMPLE_CHAT_SYSTEM_PROMPT ) +from api.chat_model import ChatCompletionRequest # Configure logging from api.logging_config import setup_logging @@ -41,59 +41,6 @@ allow_headers=["*"], # Allows all headers ) -# Models for the API -class ChatMessage(BaseModel): - role: str # 'user' or 'assistant' - content: str - -class ChatCompletionRequest(BaseModel): - """ - Model for requesting a chat completion. - """ - repo_url: str = Field(..., description="URL of the repository to query") - messages: List[ChatMessage] = Field(..., description="List of chat messages") - filePath: Optional[str] = Field(None, description="Optional path to a file in the repository to include in the prompt") - token: Optional[str] = Field(None, description="Personal access token for private repositories") - type: Optional[Literal["github", "gitlab", "bitbucket"]] = Field( - "github", - description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')", - ) - - # model parameters - provider: str = Field("google", description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)") - model: Optional[str] = Field(None, description="Model name for the specified provider") - - language: Optional[str] = Field("en", description="Language for content generation (e.g., 'en', 'ja', 'zh', 'es', 'kr', 'vi')") - excluded_dirs: List[str] = Field( - default_factory=list, - description="List or newline-separated string of directories to exclude from processing", - ) - excluded_files: List[str] = Field( - default_factory=list, - description="List or newline-separated string of file patterns to exclude from processing", - ) - included_dirs: List[str] = Field( - default_factory=list, - description="List or newline-separated string of directories to include exclusively", - ) - included_files: List[str] = Field( - default_factory=list, - description="List or newline-separated string of file patterns to include exclusively", - ) - - @field_validator( - "excluded_dirs", - "excluded_files", - "included_dirs", - "included_files", - mode="before", - ) - @classmethod - def validate_path(cls, value: list[str] | str) -> list[str]: - if isinstance(value, str): - value = [unquote(path) for path in value.strip().split("\n") if path] - return value - @app.post("/chat/completions/stream") async def chat_completions_stream(request: ChatCompletionRequest): diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index 2d7e45b71..22e8e9ca5 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -1,11 +1,8 @@ import logging from collections.abc import AsyncIterator, Callable -from typing import List, Optional -from urllib.parse import unquote from functools import partial from fastapi import WebSocket, WebSocketDisconnect -from pydantic import BaseModel, Field from api.chat import ChatStreamer, prompt_builder, is_token_limit_error from api.config import ( @@ -20,6 +17,7 @@ DEEP_RESEARCH_INTERMEDIATE_ITERATION_PROMPT, SIMPLE_CHAT_SYSTEM_PROMPT, ) +from api.chat_model import ChatCompletionRequest # Configure logging from api.logging_config import setup_logging @@ -28,34 +26,6 @@ logger = logging.getLogger(__name__) -# Models for the API -class ChatMessage(BaseModel): - role: str # 'user' or 'assistant' - content: str - -class ChatCompletionRequest(BaseModel): - """ - Model for requesting a chat completion. - """ - repo_url: str = Field(..., description="URL of the repository to query") - messages: List[ChatMessage] = Field(..., description="List of chat messages") - filePath: Optional[str] = Field(None, description="Optional path to a file in the repository to include in the prompt") - token: Optional[str] = Field(None, description="Personal access token for private repositories") - type: Optional[str] = Field("github", description="Type of repository (e.g., 'github', 'gitlab', 'bitbucket')") - - # model parameters - provider: str = Field( - "google", - description="Model provider (google, openai, openrouter, ollama, bedrock, azure, dashscope)", - ) - model: Optional[str] = Field(None, description="Model name for the specified provider") - - language: Optional[str] = Field("en", description="Language for content generation (e.g., 'en', 'ja', 'zh', 'es', 'kr', 'vi')") - excluded_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to exclude from processing") - excluded_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to exclude from processing") - included_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to include exclusively") - included_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to include exclusively") - async def handle_websocket_chat(websocket: WebSocket): """ Handle WebSocket connection for chat completions. @@ -84,25 +54,24 @@ async def handle_websocket_chat(websocket: WebSocket): request_rag = RAG(provider=request.provider, model=request.model) # Extract custom file filter parameters if provided - excluded_dirs = None - excluded_files = None - included_dirs = None - included_files = None - if request.excluded_dirs: - excluded_dirs = [unquote(dir_path) for dir_path in request.excluded_dirs.split('\n') if dir_path.strip()] - logger.info(f"Using custom excluded directories: {excluded_dirs}") + logger.info(f"Using custom excluded directories: {request.excluded_dirs}") if request.excluded_files: - excluded_files = [unquote(file_pattern) for file_pattern in request.excluded_files.split('\n') if file_pattern.strip()] - logger.info(f"Using custom excluded files: {excluded_files}") + logger.info(f"Using custom excluded files: {request.excluded_files}") if request.included_dirs: - included_dirs = [unquote(dir_path) for dir_path in request.included_dirs.split('\n') if dir_path.strip()] - logger.info(f"Using custom included directories: {included_dirs}") + logger.info(f"Using custom included directories: {request.included_dirs}") if request.included_files: - included_files = [unquote(file_pattern) for file_pattern in request.included_files.split('\n') if file_pattern.strip()] - logger.info(f"Using custom included files: {included_files}") - - request_rag.prepare_retriever(request.repo_url, request.type, request.token, excluded_dirs, excluded_files, included_dirs, included_files) + logger.info(f"Using custom included files: {request.included_files}") + + request_rag.prepare_retriever( + request.repo_url, + request.type, + request.token, + excluded_dirs=request.excluded_dirs, + excluded_files=request.excluded_files, + included_dirs=request.included_dirs, + included_files=request.included_files, + ) logger.info(f"Retriever prepared for {request.repo_url}") except ValueError as e: if "No valid documents with embeddings found" in str(e): From 95d00b436f4b8f56e5b53dfb3bfed7750cde96ee Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 11:50:54 +0800 Subject: [PATCH 21/25] import cleanup --- api/simple_chat.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/api/simple_chat.py b/api/simple_chat.py index dec70345a..a1abe5786 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -1,6 +1,5 @@ import logging -from typing import List, Optional, Callable, Literal -from urllib.parse import unquote +from typing import Callable from functools import partial from fastapi import FastAPI, HTTPException From b3361426cc393e969ae434704501002f50f92703 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 13:45:48 +0800 Subject: [PATCH 22/25] import documentations --- api/clients/__init__.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/api/clients/__init__.py b/api/clients/__init__.py index 720a1e4d8..f450ef898 100644 --- a/api/clients/__init__.py +++ b/api/clients/__init__.py @@ -1,10 +1,18 @@ +"""Unify Adalflow compatible clients to here. +Any patches and additional clients could be applied or imported in this module. +""" + from .google_embedder import GoogleEmbedderClient from .bedrock import BedrockClient from .dashscope import DashscopeClient from .ollama import OllamaClient from .litellm import LiteLLMClient from .openrouter import OpenRouterClient -from adalflow.components.model_client import AzureAIClient, OpenAIClient, GoogleGenAIClient +from adalflow.components.model_client import ( + AzureAIClient, + OpenAIClient, + GoogleGenAIClient, +) __all__ = [ "AzureAIClient", From 5de7d928c97e1fdd8c3ab99b79a4716298a40b42 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 13:59:36 +0800 Subject: [PATCH 23/25] add EOF --- api/clients/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/clients/__init__.py b/api/clients/__init__.py index f450ef898..b7ae09727 100644 --- a/api/clients/__init__.py +++ b/api/clients/__init__.py @@ -24,4 +24,4 @@ "OllamaClient", "OpenAIClient", "OpenRouterClient", -] \ No newline at end of file +] From a9e1e416e3d4e52ead4a43563da88aa7ad7174ef Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 14:16:01 +0800 Subject: [PATCH 24/25] change azureai client import path in `_stream` module. --- api/chat/_stream.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/chat/_stream.py b/api/chat/_stream.py index 5e11bb871..14525769f 100644 --- a/api/chat/_stream.py +++ b/api/chat/_stream.py @@ -198,7 +198,7 @@ class AzureChatStreamer(_OpenAICompatStreamer): ) def _build_client(self): - from adalflow.components.model_client import AzureAIClient + from api.clients import AzureAIClient return AzureAIClient() From 9ca0ae519c36f2a7074b5c5ed4c462bcb613ca92 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 14:23:48 +0800 Subject: [PATCH 25/25] follwing pep8. --- api/clients/__init__.py | 12 ++++++------ api/config.py | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/api/clients/__init__.py b/api/clients/__init__.py index b7ae09727..59661eba1 100644 --- a/api/clients/__init__.py +++ b/api/clients/__init__.py @@ -2,17 +2,17 @@ Any patches and additional clients could be applied or imported in this module. """ -from .google_embedder import GoogleEmbedderClient -from .bedrock import BedrockClient -from .dashscope import DashscopeClient -from .ollama import OllamaClient -from .litellm import LiteLLMClient -from .openrouter import OpenRouterClient from adalflow.components.model_client import ( AzureAIClient, OpenAIClient, GoogleGenAIClient, ) +from .bedrock import BedrockClient +from .dashscope import DashscopeClient +from .google_embedder import GoogleEmbedderClient +from .litellm import LiteLLMClient +from .ollama import OllamaClient +from .openrouter import OpenRouterClient __all__ = [ "AzureAIClient", diff --git a/api/config.py b/api/config.py index be11c368a..bd55b7247 100644 --- a/api/config.py +++ b/api/config.py @@ -8,15 +8,15 @@ logger = logging.getLogger(__name__) from api.clients import ( - LiteLLMClient, - OpenRouterClient, - BedrockClient, - GoogleGenAIClient, - GoogleEmbedderClient, AzureAIClient, + BedrockClient, DashscopeClient, - OpenAIClient, + GoogleEmbedderClient, + GoogleGenAIClient, + LiteLLMClient, OllamaClient, + OpenAIClient, + OpenRouterClient, ) # Get API keys from environment variables