From a9a7b25103d16e1b4d99f127397842ae2a7eac58 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Sat, 18 Jul 2026 12:01:34 +0800 Subject: [PATCH 01/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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/90] 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 From 68d990d223b5bcefa5d56729ce3cd60abf5da950 Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 16:29:16 +0800 Subject: [PATCH 26/90] using specific `mode` field to determine `deep_research` mode or normal conversation mode. --- api/chat_model.py | 1 + api/simple_chat.py | 24 ++++++++++++------------ api/websocket_wiki.py | 24 ++++++++++++------------ src/components/Ask.tsx | 19 +++++++++++++++---- src/utils/websocketClient.ts | 1 + 5 files changed, 41 insertions(+), 28 deletions(-) diff --git a/api/chat_model.py b/api/chat_model.py index f7cbf2fc4..40921212c 100644 --- a/api/chat_model.py +++ b/api/chat_model.py @@ -7,6 +7,7 @@ class ChatMessage(BaseModel): role: str # 'user' or 'assistant' content: str + mode: Literal["normal", "deep_research"] = Field(default="normal") class ChatCompletionRequest(BaseModel): """ diff --git a/api/simple_chat.py b/api/simple_chat.py index 7db123ec1..435f00dfa 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -116,21 +116,21 @@ async def chat_completions_stream(request: ChatCompletionRequest): ) # Check if this is a Deep Research request - is_deep_research = False + is_deep_research = last_message.mode == "deep_research" research_iteration = 1 - # Process messages to detect Deep Research requests - for msg in request.messages: - if hasattr(msg, 'content') and msg.content and "[DEEP RESEARCH]" in msg.content: - is_deep_research = True - # Only remove the tag from the last message - if msg == request.messages[-1]: - # Remove the Deep Research tag - msg.content = msg.content.replace("[DEEP RESEARCH]", "").strip() - # Count research iterations if this is a Deep Research request if is_deep_research: - research_iteration = sum(1 for msg in request.messages if msg.role == 'assistant') + 1 + # only count the assistant turns within the current deep research streak, so that mode switched + # before wouldn't inflate the current research iteration + current_streak = 0 + for msg in reversed(request.messages): + if msg.role == "user" and msg.mode != "deep_research": + break + if msg.role == "assistant": + current_streak += 1 + + research_iteration = current_streak + 1 logger.info(f"Deep Research request detected - iteration {research_iteration}") # Check if this is a continuation request @@ -139,7 +139,7 @@ async def chat_completions_stream(request: ChatCompletionRequest): original_topic = None for msg in request.messages: if msg.role == "user" and "continue" not in msg.content.lower(): - original_topic = msg.content.replace("[DEEP RESEARCH]", "").strip() + original_topic = msg.content.strip() logger.info(f"Found original research topic: {original_topic}") break diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index 32e37f7d9..65bd029a1 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -119,21 +119,21 @@ async def handle_websocket_chat(websocket: WebSocket): ) # Check if this is a Deep Research request - is_deep_research = False + is_deep_research = last_message.mode == "deep_research" research_iteration = 1 - # Process messages to detect Deep Research requests - for msg in request.messages: - if hasattr(msg, 'content') and msg.content and "[DEEP RESEARCH]" in msg.content: - is_deep_research = True - # Only remove the tag from the last message - if msg == request.messages[-1]: - # Remove the Deep Research tag - msg.content = msg.content.replace("[DEEP RESEARCH]", "").strip() - # Count research iterations if this is a Deep Research request if is_deep_research: - research_iteration = sum(1 for msg in request.messages if msg.role == 'assistant') + 1 + # only count the assistant turns within the current deep research streak, so that mode switched + # before wouldn't inflate the current research iteration + current_streak = 0 + for msg in reversed(request.messages): + if msg.role == "user" and msg.mode != "deep_research": + break + if msg.role == "assistant": + current_streak += 1 + + research_iteration = current_streak + 1 logger.info(f"Deep Research request detected - iteration {research_iteration}") # Check if this is a continuation request @@ -142,7 +142,7 @@ async def handle_websocket_chat(websocket: WebSocket): original_topic = None for msg in request.messages: if msg.role == "user" and "continue" not in msg.content.lower(): - original_topic = msg.content.replace("[DEEP RESEARCH]", "").strip() + original_topic = msg.content.strip() logger.info(f"Found original research topic: {original_topic}") break diff --git a/src/components/Ask.tsx b/src/components/Ask.tsx index 056afcd23..cb0f8c5b3 100644 --- a/src/components/Ask.tsx +++ b/src/components/Ask.tsx @@ -24,6 +24,7 @@ interface Provider { interface Message { role: 'user' | 'assistant' | 'system'; content: string; + mode?: 'normal' | 'deep_research'; } interface ResearchStage { @@ -299,7 +300,8 @@ const Ask: React.FC = ({ }, { role: 'user', - content: '[DEEP RESEARCH] Continue the research' + content: 'Continue the research', + mode: 'deep_research' } ]; @@ -317,7 +319,11 @@ const Ask: React.FC = ({ const requestBody: ChatCompletionRequest = { repo_url: getRepoUrl(repoInfo), type: repoInfo.type, - messages: newHistory.map(msg => ({ role: msg.role as 'user' | 'assistant', content: msg.content })), + messages: newHistory.map(msg => ({ + role: msg.role as 'user' | 'assistant', + content: msg.content, + mode: msg.mode ?? 'normal', + })), provider: selectedProvider, model: isCustomSelectedModel ? customSelectedModel : selectedModel, language: language @@ -548,7 +554,8 @@ const Ask: React.FC = ({ // Create initial message const initialMessage: Message = { role: 'user', - content: deepResearch ? `[DEEP RESEARCH] ${question}` : question + content: question, + mode: deepResearch ? 'deep_research' : 'normal' }; // Set initial conversation history @@ -559,7 +566,11 @@ const Ask: React.FC = ({ const requestBody: ChatCompletionRequest = { repo_url: getRepoUrl(repoInfo), type: repoInfo.type, - messages: newHistory.map(msg => ({ role: msg.role as 'user' | 'assistant', content: msg.content })), + messages: newHistory.map(msg => ({ + role: msg.role as 'user' | 'assistant', + content: msg.content, + mode: msg.mode ?? 'normal' + })), provider: selectedProvider, model: isCustomSelectedModel ? customSelectedModel : selectedModel, language: language diff --git a/src/utils/websocketClient.ts b/src/utils/websocketClient.ts index b29c1b104..07ffa0a39 100644 --- a/src/utils/websocketClient.ts +++ b/src/utils/websocketClient.ts @@ -17,6 +17,7 @@ const getWebSocketUrl = () => { export interface ChatMessage { role: 'user' | 'assistant' | 'system'; content: string; + mode?: 'normal' | 'deep_research'; } export interface ChatCompletionRequest { From 3e419004fc1154607526f59a1558a01ec87f277a Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Mon, 20 Jul 2026 17:32:50 +0800 Subject: [PATCH 27/90] add explicit `research_iteration` from the frontend. --- api/chat_model.py | 7 +++++++ api/simple_chat.py | 19 ++++--------------- api/websocket_wiki.py | 19 ++++--------------- src/components/Ask.tsx | 6 ++++-- src/utils/websocketClient.ts | 1 + 5 files changed, 20 insertions(+), 32 deletions(-) diff --git a/api/chat_model.py b/api/chat_model.py index 40921212c..1d3c17b05 100644 --- a/api/chat_model.py +++ b/api/chat_model.py @@ -27,6 +27,13 @@ class ChatCompletionRequest(BaseModel): 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')") + + research_iteration: int = Field( + default=1, + ge=1, + description="Current deep research iteration (1-based). Only used when the request is in deep_research mode.", + ) + excluded_dirs: List[str] = Field( default_factory=list, description="List or newline-separated string of directories to exclude from processing", diff --git a/api/simple_chat.py b/api/simple_chat.py index 435f00dfa..9c4ae413a 100644 --- a/api/simple_chat.py +++ b/api/simple_chat.py @@ -117,21 +117,10 @@ async def chat_completions_stream(request: ChatCompletionRequest): # Check if this is a Deep Research request is_deep_research = last_message.mode == "deep_research" - research_iteration = 1 # Count research iterations if this is a Deep Research request if is_deep_research: - # only count the assistant turns within the current deep research streak, so that mode switched - # before wouldn't inflate the current research iteration - current_streak = 0 - for msg in reversed(request.messages): - if msg.role == "user" and msg.mode != "deep_research": - break - if msg.role == "assistant": - current_streak += 1 - - research_iteration = current_streak + 1 - logger.info(f"Deep Research request detected - iteration {research_iteration}") + logger.info("Deep Research request detected - iteration %d",request.research_iteration) # Check if this is a continuation request if "continue" in last_message.content.lower() and "research" in last_message.content.lower(): @@ -219,10 +208,10 @@ async def chat_completions_stream(request: ChatCompletionRequest): # Create system prompt if is_deep_research: # Check if this is the first iteration - is_first_iteration = research_iteration == 1 + is_first_iteration = request.research_iteration == 1 # Check if this is the final iteration - is_final_iteration = research_iteration >= 5 + is_final_iteration = request.research_iteration >= 5 if is_first_iteration: system_prompt = DEEP_RESEARCH_FIRST_ITERATION_PROMPT.format( @@ -243,7 +232,7 @@ async def chat_completions_stream(request: ChatCompletionRequest): repo_type=repo_type, repo_url=repo_url, repo_name=repo_name, - research_iteration=research_iteration, + research_iteration=request.research_iteration, language_name=language_name ) else: diff --git a/api/websocket_wiki.py b/api/websocket_wiki.py index 65bd029a1..d7ddfda2c 100644 --- a/api/websocket_wiki.py +++ b/api/websocket_wiki.py @@ -120,21 +120,10 @@ async def handle_websocket_chat(websocket: WebSocket): # Check if this is a Deep Research request is_deep_research = last_message.mode == "deep_research" - research_iteration = 1 # Count research iterations if this is a Deep Research request if is_deep_research: - # only count the assistant turns within the current deep research streak, so that mode switched - # before wouldn't inflate the current research iteration - current_streak = 0 - for msg in reversed(request.messages): - if msg.role == "user" and msg.mode != "deep_research": - break - if msg.role == "assistant": - current_streak += 1 - - research_iteration = current_streak + 1 - logger.info(f"Deep Research request detected - iteration {research_iteration}") + logger.info("Deep Research request detected - iteration %d", request.research_iteration) # Check if this is a continuation request if "continue" in last_message.content.lower() and "research" in last_message.content.lower(): @@ -222,10 +211,10 @@ async def handle_websocket_chat(websocket: WebSocket): # Create system prompt if is_deep_research: # Check if this is the first iteration - is_first_iteration = research_iteration == 1 + is_first_iteration = request.research_iteration == 1 # Check if this is the final iteration - is_final_iteration = research_iteration >= 5 + is_final_iteration = request.research_iteration >= 5 if is_first_iteration: system_prompt = DEEP_RESEARCH_FIRST_ITERATION_PROMPT.format( @@ -246,7 +235,7 @@ async def handle_websocket_chat(websocket: WebSocket): repo_type=repo_type, repo_url=repo_url, repo_name=repo_name, - research_iteration=research_iteration, + research_iteration=request.research_iteration, language_name=language_name ) else: diff --git a/src/components/Ask.tsx b/src/components/Ask.tsx index cb0f8c5b3..e91f0b066 100644 --- a/src/components/Ask.tsx +++ b/src/components/Ask.tsx @@ -326,7 +326,8 @@ const Ask: React.FC = ({ })), provider: selectedProvider, model: isCustomSelectedModel ? customSelectedModel : selectedModel, - language: language + language: language, + research_iteration: newIteration }; // Add tokens if available @@ -573,7 +574,8 @@ const Ask: React.FC = ({ })), provider: selectedProvider, model: isCustomSelectedModel ? customSelectedModel : selectedModel, - language: language + language: language, + research_iteration: deepResearch ? 1 : undefined }; // Add tokens if available diff --git a/src/utils/websocketClient.ts b/src/utils/websocketClient.ts index 07ffa0a39..f3b22a8a1 100644 --- a/src/utils/websocketClient.ts +++ b/src/utils/websocketClient.ts @@ -29,6 +29,7 @@ export interface ChatCompletionRequest { provider?: string; model?: string; language?: string; + research_iteration?: number; excluded_dirs?: string; excluded_files?: string; } From 4e3a52904cf4837e290571cf76cef2e2d75b916a Mon Sep 17 00:00:00 2001 From: GdoongMathew Date: Tue, 21 Jul 2026 16:41:43 +0800 Subject: [PATCH 28/90] add conversation history rendering in frontend. (claude code generated) --- src/components/Ask.tsx | 580 ++++++++++++++++++++++++++++------------- 1 file changed, 403 insertions(+), 177 deletions(-) diff --git a/src/components/Ask.tsx b/src/components/Ask.tsx index e91f0b066..11eb016ee 100644 --- a/src/components/Ask.tsx +++ b/src/components/Ask.tsx @@ -34,6 +34,19 @@ interface ResearchStage { type: 'plan' | 'update' | 'conclusion'; } +// A finished exchange kept in the conversation log so previous +// question/answer pairs are preserved instead of being overwritten. +interface ConversationTurn { + question: string; + response: string; + mode: 'normal' | 'deep_research'; + // For deep research turns, every iteration/stage is preserved so the + // committed section can still browse the full multi-turn investigation. + researchStages?: ResearchStage[]; + // Index of the stage currently shown for this turn. + stageIndex?: number; +} + interface AskProps { repoInfo: RepoInfo; provider?: string; @@ -75,8 +88,13 @@ const Ask: React.FC = ({ const [conversationHistory, setConversationHistory] = useState([]); const [researchIteration, setResearchIteration] = useState(0); const [researchComplete, setResearchComplete] = useState(false); + // Completed exchanges rendered as stacked sections above the live response. + const [conversationTurns, setConversationTurns] = useState([]); + // The question that produced the currently-streaming response. + const [currentQuestion, setCurrentQuestion] = useState(''); const inputRef = useRef(null); const responseRef = useRef(null); + const conversationEndRef = useRef(null); const providerRef = useRef(provider); const modelRef = useRef(model); @@ -94,12 +112,15 @@ const Ask: React.FC = ({ } }, [onRef]); - // Scroll to bottom of response when it changes + // Keep the newest content in view as the conversation grows or streams. useEffect(() => { if (responseRef.current) { responseRef.current.scrollTop = responseRef.current.scrollHeight; } - }, [response]); + if (conversationEndRef.current) { + conversationEndRef.current.scrollIntoView({ behavior: 'smooth', block: 'end' }); + } + }, [response, conversationTurns]); // Close WebSocket when component unmounts useEffect(() => { @@ -153,6 +174,8 @@ const Ask: React.FC = ({ setQuestion(''); setResponse(''); setConversationHistory([]); + setConversationTurns([]); + setCurrentQuestion(''); setResearchIteration(0); setResearchComplete(false); setResearchStages([]); @@ -161,8 +184,16 @@ const Ask: React.FC = ({ inputRef.current.focus(); } }; - const downloadresponse = () =>{ - const blob = new Blob([response], { type: 'text/markdown' }); + + // Build the message list for previously completed turns so follow-up + // questions carry the earlier conversation as context. + const buildHistoryFromTurns = (): Message[] => + conversationTurns.flatMap(turn => [ + { role: 'user' as const, content: turn.question, mode: turn.mode }, + { role: 'assistant' as const, content: turn.response }, + ]); + const downloadresponse = (content: string = response) =>{ + const blob = new Blob([content], { type: 'text/markdown' }); const url = URL.createObjectURL(blob); const a = document.createElement('a'); a.href = url; @@ -277,6 +308,9 @@ const Ask: React.FC = ({ // WebSocket reference const webSocketRef = useRef(null); + // True while an HTTP fallback is running, so the WebSocket close handler + // doesn't also commit a (partial/errored) turn. + const fallbackActiveRef = useRef(false); // Function to continue research automatically const continueResearch = async () => { @@ -378,10 +412,12 @@ const Ask: React.FC = ({ setResponse(prev => prev + '\n\nError: WebSocket connection failed. Falling back to HTTP...'); // Fallback to HTTP if WebSocket fails - fallbackToHttp(requestBody); + fallbackActiveRef.current = true; + fallbackToHttp(requestBody, currentQuestion, 'deep_research'); }, // Close handler () => { + if (fallbackActiveRef.current) return; // Check if research is complete when the WebSocket closes const isComplete = checkIfResearchComplete(fullResponse); @@ -410,7 +446,11 @@ const Ask: React.FC = ({ }; // Fallback to HTTP if WebSocket fails - const fallbackToHttp = async (requestBody: ChatCompletionRequest) => { + const fallbackToHttp = async ( + requestBody: ChatCompletionRequest, + committedQuestion: string, + mode: 'normal' | 'deep_research' + ) => { try { // Make the API call using HTTP const apiResponse = await fetch(`/api/chat/stream`, { @@ -445,7 +485,7 @@ const Ask: React.FC = ({ // Extract research stage if this is a deep research response if (deepResearch) { - const stage = extractResearchStage(fullResponse, researchIteration); + const stage = extractResearchStage(fullResponse, requestBody.research_iteration ?? 1); if (stage) { // Add the stage to the research stages setResearchStages(prev => { @@ -462,26 +502,40 @@ const Ask: React.FC = ({ } } - // Check if research is complete - const isComplete = checkIfResearchComplete(fullResponse); - - // Force completion after a maximum number of iterations (5) - const forceComplete = researchIteration >= 5; - - if (forceComplete && !isComplete) { - // If we're forcing completion, append a comprehensive conclusion to the response - const completionNote = "\n\n## Final Conclusion\nAfter multiple iterations of deep research, we've gathered significant insights about this topic. This concludes our investigation process, having reached the maximum number of research iterations. The findings presented across all iterations collectively form our comprehensive answer to the original question."; - fullResponse += completionNote; - setResponse(fullResponse); - setResearchComplete(true); + if (mode === 'normal') { + // Normal mode: append the finished answer to the conversation log. + commitTurn(committedQuestion, fullResponse, 'normal'); } else { - setResearchComplete(isComplete); + // Deep research: mirror the WebSocket close handlers so the iteration + // counter advances and the auto-continue loop keeps running even when + // every request falls back to HTTP (the WebSocket never connects). + const iter = requestBody.research_iteration ?? 1; + if (iter === 1) { + // Iteration 1 is the research plan — always advance into the + // subsequent iterations. Advancing here (after the stream finishes) + // means the next turn captures this real answer, not a stale note. + setResearchComplete(false); + setResearchIteration(1); + } else if (iter >= 5) { + // Force a conclusion once the maximum iteration count is reached. + const isComplete = checkIfResearchComplete(fullResponse); + if (!isComplete) { + const completionNote = "\n\n## Final Conclusion\nAfter multiple iterations of deep research, we've gathered significant insights about this topic. This concludes our investigation process, having reached the maximum number of research iterations. The findings presented across all iterations collectively form our comprehensive answer to the original question."; + fullResponse += completionNote; + setResponse(fullResponse); + } + setResearchComplete(true); + } else { + // Intermediate iterations: stop only if the model signalled completion. + setResearchComplete(checkIfResearchComplete(fullResponse)); + } } } catch (error) { console.error('Error during HTTP fallback:', error); setResponse(prev => prev + '\n\nError: Failed to get a response. Please try again.'); setResearchComplete(true); } finally { + fallbackActiveRef.current = false; setIsLoading(false); } }; @@ -489,7 +543,9 @@ const Ask: React.FC = ({ // Effect to continue research when response is updated useEffect(() => { if (deepResearch && response && !isLoading && !researchComplete) { - const isComplete = checkIfResearchComplete(response); + // Iteration 1 is the plan and is never "complete"; only apply the + // completion heuristic from the second iteration onward. + const isComplete = researchIteration >= 2 && checkIfResearchComplete(response); if (isComplete) { setResearchComplete(true); } else if (researchIteration > 0 && researchIteration < 5) { @@ -504,6 +560,18 @@ const Ask: React.FC = ({ // eslint-disable-next-line react-hooks/exhaustive-deps }, [response, isLoading, deepResearch, researchComplete, researchIteration]); + // When a deep research process finishes, move the final answer into the + // conversation log so it is preserved as its own section. + useEffect(() => { + if (deepResearch && researchComplete && !isLoading && response && currentQuestion) { + commitTurn(currentQuestion, response, 'deep_research', researchStages); + setResearchIteration(0); + setResearchComplete(false); + fallbackActiveRef.current = false; + } + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [researchComplete, isLoading, researchStages]); + // Effect to update research stages when the response changes useEffect(() => { if (deepResearch && response && !isLoading) { @@ -546,21 +614,32 @@ const Ask: React.FC = ({ // Handle confirm and send request const handleConfirmAsk = async () => { + // Capture the question for this turn and clear the input so the user + // can immediately type a follow-up. + const askedQuestion = question; + const askedMode: 'normal' | 'deep_research' = deepResearch ? 'deep_research' : 'normal'; + + fallbackActiveRef.current = false; setIsLoading(true); setResponse(''); + setCurrentQuestion(askedQuestion); + setQuestion(''); setResearchIteration(0); setResearchComplete(false); + setResearchStages([]); + setCurrentStageIndex(0); try { // Create initial message const initialMessage: Message = { role: 'user', - content: question, - mode: deepResearch ? 'deep_research' : 'normal' + content: askedQuestion, + mode: askedMode }; - // Set initial conversation history - const newHistory: Message[] = [initialMessage]; + // Include previously completed turns so the backend keeps context, + // then append the new question. + const newHistory: Message[] = [...buildHistoryFromTurns(), initialMessage]; setConversationHistory(newHistory); // Prepare request body @@ -612,20 +691,23 @@ const Ask: React.FC = ({ setResponse(prev => prev + '\n\nError: WebSocket connection failed. Falling back to HTTP...'); // Fallback to HTTP if WebSocket fails - fallbackToHttp(requestBody); + fallbackActiveRef.current = true; + fallbackToHttp(requestBody, askedQuestion, askedMode); }, // Close handler () => { + if (fallbackActiveRef.current) return; // If deep research is enabled, check if we should continue if (deepResearch) { - const isComplete = checkIfResearchComplete(fullResponse); - setResearchComplete(isComplete); - - // If not complete, start the research process - if (!isComplete) { - setResearchIteration(1); - // The continueResearch function will be triggered by the useEffect - } + // Iteration 1 is the initial research plan — always continue into + // the deeper iterations instead of letting the completion + // heuristic (which can false-positive on the plan text) stop here. + setResearchComplete(false); + setResearchIteration(1); + // The continueResearch function will be triggered by the useEffect. + } else { + // Normal mode: the answer is final, so append it to the log. + commitTurn(askedQuestion, fullResponse, 'normal'); } setIsLoading(false); @@ -639,6 +721,54 @@ const Ask: React.FC = ({ } }; + // Move a finished exchange out of the live response area and into the + // persistent conversation log. + const commitTurn = ( + turnQuestion: string, + turnResponse: string, + mode: 'normal' | 'deep_research', + stages: ResearchStage[] = [] + ) => { + if (!turnQuestion && !turnResponse) return; + + // For deep research, make sure the final answer is always present as the + // last stage (stage extraction can miss it, e.g. on forced completion). + let finalStages = stages; + if (mode === 'deep_research') { + const last = stages[stages.length - 1]; + if (turnResponse && (!last || last.content !== turnResponse)) { + finalStages = [ + ...stages, + { + title: 'Final Result', + content: turnResponse, + iteration: stages.length + 1, + type: 'conclusion', + }, + ]; + } + } + + setConversationTurns(prev => [...prev, { + question: turnQuestion, + response: turnResponse, + mode, + researchStages: finalStages, + stageIndex: finalStages.length > 0 ? finalStages.length - 1 : 0, + }]); + setResponse(''); + setCurrentQuestion(''); + setResearchStages([]); + setCurrentStageIndex(0); + }; + + // Switch which research stage is shown for a committed deep research turn. + const setTurnStage = (turnIndex: number, stageIndex: number) => { + setConversationTurns(prev => + prev.map((turn, i) => (i === turnIndex ? { ...turn, stageIndex } : turn)) + ); + }; + const [buttonWidth, setButtonWidth] = useState(0); const buttonRef = useRef(null); @@ -651,165 +781,153 @@ const Ask: React.FC = ({ }, [messages.ask?.askButton, isLoading]); return ( -
-
-
- {/* Model selection button */} - -
- - {/* Question input */} -
-
- setQuestion(e.target.value)} - placeholder={messages.ask?.placeholder || 'What would you like to know about this codebase?'} - className="block w-full rounded-md border border-[var(--border-color)] bg-[var(--input-bg)] text-[var(--foreground)] px-5 py-3.5 text-base shadow-sm focus:border-[var(--accent-primary)] focus:ring-2 focus:ring-[var(--accent-primary)]/30 focus:outline-none transition-all" - style={{ paddingRight: `${buttonWidth + 24}px` }} - disabled={isLoading} - /> - -
+
+ {/* Conversation log: previous turns are preserved and new results + are appended below as their own sections. */} +
+ {conversationTurns.map((turn, idx) => { + // Deep research turns keep every iteration; show the selected one + // and expose navigation. Normal turns just show their response. + const stages = turn.researchStages ?? []; + const hasStages = turn.mode === 'deep_research' && stages.length > 0; + const activeStage = hasStages + ? Math.min(turn.stageIndex ?? stages.length - 1, stages.length - 1) + : 0; + const displayContent = hasStages ? stages[activeStage].content : turn.response; + + return ( +
+ {/* User question */} +
+
+ {turn.question} +
+
- {/* Deep Research toggle */} -
-
-