From 7290e315ef3100d7bf7ef8e336154737dc179f4c Mon Sep 17 00:00:00 2001 From: Herman Haggerty Date: Mon, 1 Sep 2025 15:53:39 -0500 Subject: [PATCH 1/3] Implement Chat with RAG --- CHANGELOG.md | 7 - openwebui/api/chats.py | 247 ++++++++++++++++++-------------- openwebui/api/knowledge.py | 94 +++++++++++- openwebui/cli/main.py | 156 +++++++++++++++++--- tests/cli/test_cli_chat.py | 162 +++++++++++++++++++++ tests/sdk/test_sdk_chats.py | 222 ++++++++++++++++++++-------- tests/sdk/test_sdk_knowledge.py | 152 +++++++++++++++++++- 7 files changed, 842 insertions(+), 198 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 72d7f45..cb6a50f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,11 +4,4 @@ ## v0.1.0 (2025-08-29) -- Initial Release -## v1.0.1 (2025-08-26) - - -## v1.0.0 (2025-08-26) - -- Initial Release diff --git a/openwebui/api/chats.py b/openwebui/api/chats.py index 362e8b7..7f85219 100644 --- a/openwebui/api/chats.py +++ b/openwebui/api/chats.py @@ -1,10 +1,12 @@ import logging -from typing import List, Optional +from typing import List, Optional, Any import httpx -from ..exceptions import ConnectionError -from ..utils.api_utils import handle_api_response # NEW IMPORT for centralized handler +# Import the new KnowledgeBaseAPI to perform RAG queries +from .knowledge import KnowledgeBaseAPI +from ..exceptions import APIError, AuthenticationError, ConnectionError, NotFoundError +from ..utils.api_utils import handle_api_response # Import the generated client, models, and specific api functions from ..open_web_ui_client.open_web_ui_client import AuthenticatedClient, models @@ -19,8 +21,7 @@ generate_chat_completion_openai_chat_completions_post, ) from ..open_web_ui_client.open_web_ui_client.api.folders import get_folder_by_id_api_v1_folders_id_get -from ..open_web_ui_client.open_web_ui_client.types import UNSET # Still needed for function signature - +from ..open_web_ui_client.open_web_ui_client.types import UNSET log = logging.getLogger(__name__) @@ -30,33 +31,82 @@ class ChatsAPI: def __init__(self, client: AuthenticatedClient): self._client = client - - async def create(self, model: str, prompt: str, folder_id: Optional[str] = None) -> models.ChatResponse: - """Creates a new chat session by first getting an LLM response and then saving the full conversation.""" + # Instantiate KnowledgeBaseAPI to be used for RAG operations + self.knowledge = KnowledgeBaseAPI(client) + + async def create( + self, + model: str, + prompt: str, + folder_id: Optional[str] = None, + kb_ids: Optional[List[str]] = None, + k: Optional[int] = None, + k_reranker: Optional[int] = None, + r: Optional[float] = None, + hybrid: Optional[bool] = None, + hybrid_bm25_weight: Optional[float] = None, + ) -> models.ChatResponse: + """ + Creates a new chat session. If kb_ids are provided, it performs a RAG query + to augment the prompt with context from the knowledge base(s). + """ log.info(f"Creating new chat with model '{model}'.") + + final_prompt = prompt + + # --- RAG WORKFLOW ADDITION --- + if kb_ids: + print(kb_ids) + log.info(f"RAG workflow activated. Querying KBs {kb_ids} for context.") + try: + retrieved_chunks = await self.knowledge.query( + prompt, + kb_ids, + k=k, + k_reranker=k_reranker, + r=r, + hybrid=hybrid, + hybrid_bm25_weight=hybrid_bm25_weight + ) + if retrieved_chunks: + # Combine the retrieved chunks into a single context string. + # This assumes each chunk is a dict with a 'content' field. + retrieved_context = "\n\n---\n\n".join( + [chunk.get("content", "") for chunk in retrieved_chunks if isinstance(chunk, dict)]) + + if retrieved_context: # Ensure there's actually content before augmenting + final_prompt = ( + f"Please use the following context to answer the question.\n\n" + f"--- Context ---\n{retrieved_context}\n\n" + f"--- Question ---\n{prompt}" + ) + log.debug("Prompt augmented with retrieved context from KB(s).") + else: + log.warning(f"Retrieved chunks from KB(s) {kb_ids} had no usable content.") + else: + log.warning(f"No chunks retrieved from KB(s) {kb_ids} for the prompt.") + except Exception as e: + log.error(f"Failed to query knowledge base(s) {kb_ids}. Proceeding without RAG. Error: {e}") + # --- END RAG WORKFLOW --- + try: log.debug("Step 1/2: Getting LLM completion.") completion_payload = models.GenerateChatCompletionOpenaiChatCompletionsPostFormData() completion_payload.additional_properties = { "model": model, - "messages": [{"role": "user", "content": prompt}], + "messages": [{"role": "user", "content": final_prompt}], "stream": False, } - completion_response = ( - await generate_chat_completion_openai_chat_completions_post.asyncio_detailed( - client=self._client, body=completion_payload - ) + completion_response = await generate_chat_completion_openai_chat_completions_post.asyncio_detailed( + client=self._client, body=completion_payload ) - # Use the centralized handler completion_data = handle_api_response(completion_response, "LLM completion") - assistant_content = completion_data["choices"][0]["message"][ - "content" - ] # Assuming completion_data is a dict here + assistant_content = completion_data["choices"][0]["message"]["content"] log.debug("LLM completion received successfully.") - # Step 2: Create the chat session with the full conversation log.debug("Step 2/2: Saving full conversation to a new chat session.") + # Important: Store the original prompt in the chat history, not the augmented one. full_messages = [ {"role": "user", "content": prompt}, {"role": "assistant", "content": assistant_content}, @@ -68,7 +118,6 @@ async def create(self, model: str, prompt: str, folder_id: Optional[str] = None) create_response = await create_new_chat_api_v1_chats_new_post.asyncio_detailed( client=self._client, body=create_payload ) - # Use the centralized handler new_chat = handle_api_response(create_response, "chat creation") log.info(f"Successfully created new chat with ID: {new_chat.id}") return new_chat @@ -76,37 +125,77 @@ async def create(self, model: str, prompt: str, folder_id: Optional[str] = None) except httpx.ConnectError as e: raise ConnectionError(f"A network error occurred: {e}") from e - async def continue_chat(self, chat_id: str, prompt: str) -> models.ChatResponse: - """Adds a new message to an existing chat and gets a response from the LLM.""" + async def continue_chat( + self, + chat_id: str, + prompt: str, + kb_ids: Optional[List[str]] = None, + k: Optional[int] = None, + k_reranker: Optional[int] = None, + r: Optional[float] = None, + hybrid: Optional[bool] = None, + hybrid_bm25_weight: Optional[float] = None, + ) -> models.ChatResponse: + """ + Adds a new message to an existing chat. If kb_ids are provided, it performs a RAG query + to augment the prompt with context from the knowledge base(s). + """ log.info(f"Continuing chat with ID: {chat_id}") + + final_prompt = prompt + + # --- RAG WORKFLOW ADDITION --- + if kb_ids: + log.info(f"RAG workflow activated for continuation. Querying KBs {kb_ids} for context.") + try: + retrieved_chunks = await self.knowledge.query( + prompt, + kb_ids, + k=k, + k_reranker=k_reranker, + r=r, + hybrid=hybrid, + hybrid_bm25_weight=hybrid_bm25_weight + ) + if retrieved_chunks: + retrieved_context = "\n\n---\n\n".join( + [chunk.get("content", "") for chunk in retrieved_chunks if isinstance(chunk, dict)]) + if retrieved_context: + final_prompt = ( + f"Please use the following context to answer the question.\n\n" + f"--- Context ---\n{retrieved_context}\n\n" + f"--- Question ---\n{prompt}" + ) + log.debug("Follow-up prompt augmented with retrieved context.") + else: + log.warning(f"Retrieved chunks from KB(s) {kb_ids} had no usable content for follow-up.") + else: + log.warning(f"No chunks retrieved from KB(s) {kb_ids} for the follow-up prompt.") + except Exception as e: + log.error(f"Failed to query knowledge base(s) {kb_ids}. Proceeding without RAG. Error: {e}") + # --- END RAG WORKFLOW --- + try: log.debug(f"Step 1/3: Fetching history for chat_id='{chat_id}'.") - chat_details = await self.get(chat_id) # Reuse existing get method + chat_details = await self.get(chat_id) - model = chat_details.chat.additional_properties.get("models")[ - 0 - ] # Assumes chat_details is models.ChatResponse - messages = chat_details.chat.additional_properties.get( - "messages", [] - ) # Assumes chat_details is models.ChatResponse - messages.append({"role": "user", "content": prompt}) + model = chat_details.chat.additional_properties.get("models")[0] + messages = chat_details.chat.additional_properties.get("messages", []) + messages.append({"role": "user", "content": final_prompt}) log.debug(f"Step 2/3: Sending continued prompt to model '{model}'.") completion_payload = models.GenerateChatCompletionOpenaiChatCompletionsPostFormData() completion_payload.additional_properties = {"model": model, "messages": messages, "stream": False} - completion_response = ( - await generate_chat_completion_openai_chat_completions_post.asyncio_detailed( - client=self._client, body=completion_payload - ) + completion_response = await generate_chat_completion_openai_chat_completions_post.asyncio_detailed( + client=self._client, body=completion_payload ) - # Use the centralized handler completion_data = handle_api_response(completion_response, "LLM completion for continue chat") - assistant_response = completion_data["choices"][0]["message"][ - "content" - ] # Assuming completion_data is a dict here + assistant_response = completion_data["choices"][0]["message"]["content"] log.debug(f"Step 3/3: Appending assistant's response and updating chat '{chat_id}' on server.") + # Important: Replace the augmented prompt with the original for a clean history + messages[-1] = {"role": "user", "content": prompt} messages.append({"role": "assistant", "content": assistant_response}) updated_chat_data = models.ChatFormChat() @@ -122,26 +211,19 @@ async def continue_chat(self, chat_id: str, prompt: str) -> models.ChatResponse: id=chat_id, client=self._client, body=update_payload ) - # Use the centralized handler updated_chat = handle_api_response(update_response, "chat update") log.info(f"Successfully updated chat with ID: {chat_id}") return updated_chat except httpx.ConnectError as e: raise ConnectionError(f"A network error occurred: {e}") from e - async def list(self) -> List[models.ChatTitleIdResponse]: """ Retrieves a list of all chat titles and IDs for the authenticated user. - Returns: - A list of ChatTitleIdResponse objects. """ log.info("Listing all chats for the user.") try: - response = await get_session_user_chat_list_api_v1_chats_list_get.asyncio_detailed( - client=self._client - ) - # Use the centralized handler + response = await get_session_user_chat_list_api_v1_chats_list_get.asyncio_detailed(client=self._client) return handle_api_response(response, "chat list") except httpx.ConnectError as e: raise ConnectionError(f"A network error occurred: {e}") from e @@ -149,52 +231,30 @@ async def list(self) -> List[models.ChatTitleIdResponse]: async def get(self, chat_id: str) -> models.ChatResponse: """ Fetches the full details and message history of a specific chat. - Returns: - A ChatResponse object containing chat details. """ log.info(f"Getting details for chat_id: {chat_id}") try: - response = await get_chat_by_id_api_v1_chats_id_get.asyncio_detailed( - id=chat_id, client=self._client - ) - # Use the centralized handler + response = await get_chat_by_id_api_v1_chats_id_get.asyncio_detailed(id=chat_id, client=self._client) return handle_api_response(response, f"chat details for ID {chat_id}") except httpx.ConnectError as e: raise ConnectionError(f"A network error occurred: {e}") from e async def delete(self, chat_id: str) -> bool: - """ - Deletes a chat by its ID. - Returns: - True if the deletion was successful. - """ + """Deletes a chat by its ID.""" log.info(f"Deleting chat with ID: {chat_id}") try: - response = await delete_chat_by_id_api_v1_chats_id_delete.asyncio_detailed( - id=chat_id, client=self._client - ) - # handle_api_response will return True for successful deletes + response = await delete_chat_by_id_api_v1_chats_id_delete.asyncio_detailed(id=chat_id, client=self._client) return handle_api_response(response, f"chat deletion for ID {chat_id}") except httpx.ConnectError as e: raise ConnectionError(f"A network error occurred: {e}") from e async def rename(self, chat_id: str, new_title: str) -> models.ChatResponse: - """ - Renames (updates the title of) an existing chat. - Returns: - A ChatResponse object representing the updated chat. - """ + """Renames (updates the title of) an existing chat.""" log.info(f"Renaming chat '{chat_id}' to '{new_title}'.") try: - # Step 1: Get the existing chat details (full object for update) - log.debug(f"Fetching chat details for renaming: {chat_id}") - chat_details = await self.get(chat_id) # Reuse existing get method - - # Step 2: Update only the title + chat_details = await self.get(chat_id) chat_details.title = new_title - # Step 3: Reconstruct the ChatForm body for the update API call - # The chat content itself is within chat_details.chat.additional_properties chat_data_payload = models.ChatFormChat() chat_data_payload.additional_properties = chat_details.chat.additional_properties @@ -203,59 +263,26 @@ async def rename(self, chat_id: str, new_title: str) -> models.ChatResponse: folder_id=chat_details.folder_id if chat_details.folder_id is not UNSET else UNSET, ) - log.debug(f"Sending update request for chat: {chat_id}") response = await update_chat_by_id_api_v1_chats_id_post.asyncio_detailed( id=chat_id, client=self._client, body=update_payload ) - # Use the centralized handler updated_chat = handle_api_response(response, "chat rename") log.info(f"Successfully renamed chat '{chat_id}' to '{updated_chat.title}'.") return updated_chat - except httpx.ConnectError as e: raise ConnectionError(f"A network error occurred: {e}") from e async def list_by_folder(self, folder_id: str) -> List[models.ChatTitleIdResponse]: - """ - Lists all chat IDs and titles within a specific folder. - - Notes: - The Open WebUI API documentation suggests that `GET /api/v1/folders/{id}` - returns a `FolderModel` which contains an `items` field. - This `items` field may contain chats. - - This method retrieves folder details and extracts chats from its `items` property. - - Args: - folder_id (str): The ID of the folder to list chats from. - - Returns: - List[models.ChatTitleIdResponse]: A list of chat title and ID objects. - - Raises: - NotFoundError: If the folder with the given ID does not exist. - AuthenticationError, APIError, ConnectionError - """ + """Lists all chat IDs and titles within a specific folder.""" log.info(f"Listing chats for folder_id: {folder_id}") try: - # Use the generated folders client to get folder details - log.debug(f"Fetching folder details for ID: {folder_id} to list its chats.") - - folder_response = await get_folder_by_id_api_v1_folders_id_get.asyncio_detailed( - id=folder_id, client=self._client - ) - # Use the centralized handler + folder_response = await get_folder_by_id_api_v1_folders_id_get.asyncio_detailed(id=folder_id, + client=self._client) folder_details = handle_api_response(folder_response, f"folder details for {folder_id}") - # FIX: Check if folder_details.items is None before accessing .chats - # If items is None, there are no chats in the folder. - chats_in_folder = ( - folder_details.items.chats if folder_details.items and folder_details.items.chats else [] - ) - + chats_in_folder = folder_details.items.chats if folder_details.items and folder_details.items.chats else [] log.info(f"Found {len(chats_in_folder)} chats in folder '{folder_id}'.") return chats_in_folder - except httpx.ConnectError as e: - raise ConnectionError(f"A network error occurred: {e}") from e + raise ConnectionError(f"A network error occurred: {e}") from e \ No newline at end of file diff --git a/openwebui/api/knowledge.py b/openwebui/api/knowledge.py index a7d9a71..924dc6e 100644 --- a/openwebui/api/knowledge.py +++ b/openwebui/api/knowledge.py @@ -15,7 +15,7 @@ # --- DIRECT IMPORTS from the generated OpenWebUI API client --- from ..open_web_ui_client.open_web_ui_client import AuthenticatedClient, models -from ..open_web_ui_client.open_web_ui_client.types import File as OpenAPIGeneratedFile +from ..open_web_ui_client.open_web_ui_client.types import File as OpenAPIGeneratedFile, UNSET # API function imports (extrapolated from kbmanager.api_interface and likely OpenWebUI structure) from ..open_web_ui_client.open_web_ui_client.api.knowledge.create_new_knowledge_api_v1_knowledge_create_post import ( @@ -39,6 +39,9 @@ from ..open_web_ui_client.open_web_ui_client.api.knowledge.delete_knowledge_by_id_api_v1_knowledge_id_delete_delete import ( asyncio_detailed as delete_kb_api_call, ) +from ..open_web_ui_client.open_web_ui_client.api.retrieval.query_collection_handler_api_v1_retrieval_query_collection_post import ( + asyncio_detailed as query_kb_api_call, +) log = logging.getLogger(__name__) @@ -107,6 +110,95 @@ async def delete(self, kb_id: str) -> bool: except httpx.ConnectError as e: raise ConnectionError(f"A network error occurred while deleting knowledge base: {e}") from e + async def query( + self, + query_text: str, + kb_ids: List[str], + k: Optional[int] = None, + k_reranker: Optional[int] = None, + r: Optional[float] = None, + hybrid: Optional[bool] = None, + hybrid_bm25_weight: Optional[float] = None, + ) -> List[Dict[str, Any]]: # This return type is still correct conceptually for the final output + """ + Queries one or more knowledge bases and returns the most relevant text chunks. + + Args: + query_text: The user's query to search for. + kb_ids: A list of knowledge base IDs to query. + k: The number of results to return (top-k). + k_reranker: The number of results to re-rank (if reranking is used). + r: The relevance score threshold (0.0 to 1.0). + hybrid: Whether to use hybrid search (vector + keyword). + hybrid_bm25_weight: The weight for BM25 in hybrid search (0.0 to 1.0). + + Returns: + A list of dictionary objects, each representing a retrieved document chunk. + Each dictionary will typically contain 'content' and 'meta' fields. + """ + log.info(f"Querying KBs: {kb_ids} with text: '{query_text[:50]}...'") + try: + # Construct QueryCollectionsForm dynamically based on provided parameters + query_form_data = { + "collection_names": kb_ids, + "query": query_text, + } + if k is not None: + query_form_data["k"] = k + if k_reranker is not None: + query_form_data["k_reranker"] = k_reranker + if r is not None: + query_form_data["r"] = r + if hybrid is not None: + query_form_data["hybrid"] = hybrid + if hybrid_bm25_weight is not None: + query_form_data["hybrid_bm25_weight"] = hybrid_bm25_weight + + query_form = models.QueryCollectionsForm( + collection_names=query_form_data["collection_names"], + query=query_form_data["query"], + k=query_form_data.get("k", UNSET), + k_reranker=query_form_data.get("k_reranker", UNSET), + r=query_form_data.get("r", UNSET), + hybrid=query_form_data.get("hybrid", UNSET), + hybrid_bm25_weight=query_form_data.get("hybrid_bm25_weight", UNSET), + ) + + response = await query_kb_api_call(client=self._client, body=query_form) + retrieved_response_dict = handle_api_response(response, f"query on KBs {kb_ids}") + + # --- FIX STARTS HERE --- + if not isinstance(retrieved_response_dict, dict): + log.error( + f"Expected dict from retrieval API, but received: {type(retrieved_response_dict)}. Raw: {retrieved_response_dict}") + raise APIError(f"Unexpected retrieval API response type: {type(retrieved_response_dict)}", + response.status_code) + + # The actual chunks are in the 'documents' field, and 'metadatas' often stores the corresponding meta + # Assuming 'documents' is a list of lists, and we want to flatten it and pair with metadata. + # The structure from your error message: {'distances': [...], 'documents': [[...]], 'metadatas': [[...]]} + + extracted_documents_lists = retrieved_response_dict.get('documents', []) + extracted_metadatas_lists = retrieved_response_dict.get('metadatas', []) + + # Flatten the lists and pair documents with their metadata + # Assuming a 1:1 mapping between documents in inner list and their metadatas + retrieved_chunks = [] + for doc_list, meta_list in zip(extracted_documents_lists, extracted_metadatas_lists): + for doc_content, meta_data in zip(doc_list, meta_list): + retrieved_chunks.append({ + "content": doc_content, + "meta": meta_data # This could be any relevant metadata + }) + # --- FIX ENDS HERE --- + + log.info(f"Retrieved {len(retrieved_chunks)} chunks from KBs: {kb_ids}.") + log.debug(f"Retrieved chunks: {retrieved_chunks}") + return retrieved_chunks # Now returning a List[Dict] as expected by ChatsAPI + + except httpx.ConnectError as e: + raise ConnectionError(f"A network error occurred while querying KBs '{kb_ids}': {e}") from e + async def list_all(self) -> List[models.KnowledgeResponse]: """ Lists all available knowledge bases in OpenWebUI. diff --git a/openwebui/cli/main.py b/openwebui/cli/main.py index 8cb34fc..4279912 100644 --- a/openwebui/cli/main.py +++ b/openwebui/cli/main.py @@ -3,7 +3,7 @@ import logging import json from pathlib import Path # Added for path operations in upload_dir & update_file -from typing import Optional, Any +from typing import Optional, Any, List import click import httpx @@ -23,7 +23,7 @@ handler.setFormatter(formatter) log.addHandler(handler) - +logging.basicConfig(level=logging.DEBUG) # --- Output Formatting Helper --- def format_output(data: Any, output_format: str): @@ -97,22 +97,81 @@ def chat(): """Commands for managing chats.""" pass - +# create_chat_command_wrapper function (add new options for RAG settings) @chat.command("create") @click.argument("prompt") @click.option("--model", "-m", default="gemini-1.5-flash", help="The model name to use.") @click.option("--folder-id", help="Optional ID of the folder to add this chat to.") +@click.option( + "--kb-id", "kb_ids", + multiple=True, + help="ID of a knowledge base to use for RAG. Can be specified multiple times." +) +@click.option("--k", type=int, help="Number of top hits to retrieve from the KB.") +@click.option("--k-reranker", type=int, help="Number of re-ranked hits from the KB.") +@click.option("--r", type=float, help="Relevance score threshold for KB retrieval (0.0 to 1.0).") +@click.option("--hybrid/--no-hybrid", default=None, type=bool, help="Enable/disable hybrid search for KB retrieval.") +@click.option("--hybrid-bm25-weight", type=float, help="Weight for BM25 in hybrid search (0.0 to 1.0).") @click.pass_context -def create_chat_command_wrapper(ctx, prompt: str, model: str, folder_id: Optional[str]): - """Creates a new chat, gets a response, and saves it to Open WebUI.""" - asyncio.run(_create_chat_async(ctx, prompt, model, folder_id)) - - -async def _create_chat_async(ctx, prompt: str, model: str, folder_id: Optional[str]): +def create_chat_command_wrapper( + ctx, + prompt: str, + model: str, + folder_id: Optional[str], + kb_ids: tuple[str], + k: Optional[int], + k_reranker: Optional[int], + r: Optional[float], + hybrid: Optional[bool], + hybrid_bm25_weight: Optional[float], +): + """Creates a new chat, optionally using one or more knowledge bases with RAG.""" + kb_id_list = list(kb_ids) if kb_ids else None + asyncio.run( + _create_chat_async( + ctx, + prompt, + model, + folder_id, + kb_id_list, + k, + k_reranker, + r, + hybrid, + hybrid_bm25_weight, + ) + ) + +# _create_chat_async function (update signature to receive new RAG settings) +async def _create_chat_async( + ctx, + prompt: str, + model: str, + folder_id: Optional[str], + kb_ids: Optional[List[str]], + k: Optional[int], + k_reranker: Optional[int], + r: Optional[float], + hybrid: Optional[bool], + hybrid_bm25_weight: Optional[float], +): log.info(f"CLI: Attempting to create chat with prompt: '{prompt[:50]}...'") + if kb_ids: + log.info(f"CLI: Using knowledge bases: {kb_ids}") + try: sdk = OpenWebUI() - result_chat = await sdk.chats.create(model=model, prompt=prompt, folder_id=folder_id) + result_chat = await sdk.chats.create( + model=model, + prompt=prompt, + folder_id=folder_id, + kb_ids=kb_ids, + k=k, + k_reranker=k_reranker, + r=r, + hybrid=hybrid, + hybrid_bm25_weight=hybrid_bm25_weight, + ) if ctx.obj["OUTPUT_FORMAT"] == "json": format_output(result_chat, ctx.obj["OUTPUT_FORMAT"]) @@ -120,13 +179,11 @@ async def _create_chat_async(ctx, prompt: str, model: str, folder_id: Optional[s click.secho( f"✅ Success! New chat created with ID: {result_chat.id}", fg="bright_green", bold=True ) - # Access nested attributes defensively based on common generated client patterns last_message_content = "" - if hasattr(result_chat, "chat") and hasattr(result_chat.chat, "additional_properties"): + if hasattr(result_chat, 'chat') and hasattr(result_chat.chat, 'additional_properties'): messages = result_chat.chat.additional_properties.get("messages", []) if messages: last_message_content = messages[-1].get("content", "") - click.secho("Assistant Response:", fg="cyan", bold=True) click.echo(last_message_content) log.info("CLI: Chat creation command completed successfully.") @@ -139,27 +196,83 @@ async def _create_chat_async(ctx, prompt: str, model: str, folder_id: Optional[s raise click.Abort() +# continue_chat_command_wrapper function (add new options for RAG settings) @chat.command("continue") @click.argument("chat_id") @click.argument("prompt") +@click.option( + "--kb-id", "kb_ids", + multiple=True, + help="ID of a knowledge base to use for RAG. Can be specified multiple times." +) +@click.option("--k", type=int, help="Number of top hits to retrieve from the KB.") +@click.option("--k-reranker", type=int, help="Number of re-ranked hits from the KB.") +@click.option("--r", type=float, help="Relevance score threshold for KB retrieval (0.0 to 1.0).") +@click.option("--hybrid/--no-hybrid", default=None, type=bool, help="Enable/disable hybrid search for KB retrieval.") +@click.option("--hybrid-bm25-weight", type=float, help="Weight for BM25 in hybrid search (0.0 to 1.0).") @click.pass_context -def continue_chat_command_wrapper(ctx, chat_id: str, prompt: str): - """Continues an existing chat thread by its ID.""" - asyncio.run(_continue_chat_async(ctx, chat_id, prompt)) - +def continue_chat_command_wrapper( + ctx, + chat_id: str, + prompt: str, + kb_ids: tuple[str], + k: Optional[int], + k_reranker: Optional[int], + r: Optional[float], + hybrid: Optional[bool], + hybrid_bm25_weight: Optional[float], +): + """Continues an existing chat thread by its ID, optionally using RAG.""" + kb_id_list = list(kb_ids) if kb_ids else None + asyncio.run( + _continue_chat_async( + ctx, + chat_id, + prompt, + kb_id_list, + k, + k_reranker, + r, + hybrid, + hybrid_bm25_weight, + ) + ) + +# _continue_chat_async function (update signature to receive new RAG settings) +async def _continue_chat_async( + ctx, + chat_id: str, + prompt: str, + kb_ids: Optional[List[str]], + k: Optional[int], + k_reranker: Optional[int], + r: Optional[float], + hybrid: Optional[bool], + hybrid_bm25_weight: Optional[float], +): + log.info(f"CLI: Attempting to continue chat '{chat_id}' with prompt: '{prompt}...'") + if kb_ids: + log.info(f"CLI: Using knowledge bases: {kb_ids}") -async def _continue_chat_async(ctx, chat_id: str, prompt: str): - log.info(f"CLI: Attempting to continue chat '{chat_id}' with prompt: '{prompt[:50]}...'") try: sdk = OpenWebUI() - updated_chat = await sdk.chats.continue_chat(chat_id=chat_id, prompt=prompt) + updated_chat = await sdk.chats.continue_chat( + chat_id=chat_id, + prompt=prompt, + kb_ids=kb_ids, + k=k, + k_reranker=k_reranker, + r=r, + hybrid=hybrid, + hybrid_bm25_weight=hybrid_bm25_weight, + ) if ctx.obj["OUTPUT_FORMAT"] == "json": format_output(updated_chat, ctx.obj["OUTPUT_FORMAT"]) else: click.secho(f"✅ Success! Chat {updated_chat.id} updated.", fg="bright_green", bold=True) last_message_content = "" - if hasattr(updated_chat, "chat") and hasattr(updated_chat.chat, "additional_properties"): + if hasattr(updated_chat, 'chat') and hasattr(updated_chat.chat, 'additional_properties'): messages = updated_chat.chat.additional_properties.get("messages", []) if messages: last_message_content = messages[-1].get("content", "") @@ -174,7 +287,6 @@ async def _continue_chat_async(ctx, chat_id: str, prompt: str): log.error("CLI: Chat continuation command failed.") raise click.Abort() - @chat.command("list") @click.argument("chat_id") # This lists messages OF a specific chat @click.pass_context diff --git a/tests/cli/test_cli_chat.py b/tests/cli/test_cli_chat.py index cadc9df..c8ad290 100644 --- a/tests/cli/test_cli_chat.py +++ b/tests/cli/test_cli_chat.py @@ -104,3 +104,165 @@ def test_chat_delete_cli_aborted(runner, mock_sdk_client): assert result.exit_code == 1 assert "Aborted." in result.output mock_sdk_client.chats.delete.assert_not_awaited() + +def test_chat_create_cli_with_rag_success(runner, mock_sdk_client): + """Test `owui chat create` with RAG parameters.""" + mock_chat_response = MagicMock(spec=models.ChatResponse, id="new-chat-rag", title="RAG Test Chat") + mock_chat_response.chat.additional_properties = { + "messages": [ + {"role": "user", "content": "User prompt about data."}, + {"role": "assistant", "content": "Assistant response based on RAG."}, + ] + } + mock_sdk_client.chats.create.return_value = mock_chat_response + + result = runner.invoke( + cli, + [ + "chat", + "create", + "What happened in 2023?", + "--model", "test-rag-model", + "--kb-id", "history-kb-1", + "--kb-id", "news-kb-2", + "--k", "5", + "--k-reranker", "2", + "--r", "0.7", + "--hybrid", + "--hybrid-bm25-weight", "0.33", + "--folder-id", "my-test-folder" + ] + ) + + assert result.exit_code == 0 + assert "✅ Success! New chat created with ID: new-chat-rag" in result.output + assert "Assistant response based on RAG." in result.output + + # Verify that the SDK's chat.create method was called with the correct parameters + mock_sdk_client.chats.create.assert_awaited_once_with( + model="test-rag-model", + prompt="What happened in 2023?", + folder_id="my-test-folder", + kb_ids=["history-kb-1", "news-kb-2"], # Should be a list + k=5, + k_reranker=2, + r=0.7, + hybrid=True, + hybrid_bm25_weight=0.33, + ) + +def test_chat_create_cli_with_rag_no_hybrid_success(runner, mock_sdk_client): + """Test `owui chat create` with RAG parameters, explicitly no-hybrid.""" + mock_chat_response = MagicMock(spec=models.ChatResponse, id="new-chat-no-hybrid", title="No Hybrid RAG Chat") + mock_chat_response.chat.additional_properties = { + "messages": [ + {"role": "user", "content": "Tell me something."}, + {"role": "assistant", "content": "No hybrid response."}, + ] + } + mock_sdk_client.chats.create.return_value = mock_chat_response + + result = runner.invoke( + cli, + [ + "chat", + "create", + "Tell me about the universe.", + "--kb-id", "cosmos-kb", + "--no-hybrid", # Explicitly set no hybrid + ] + ) + + assert result.exit_code == 0 + assert "✅ Success! New chat created with ID: new-chat-no-hybrid" in result.output + + mock_sdk_client.chats.create.assert_awaited_once_with( + model="gemini-1.5-flash", # Default model + prompt="Tell me about the universe.", + folder_id=None, + kb_ids=["cosmos-kb"], + k=None, + k_reranker=None, + r=None, + hybrid=False, # Should be False + hybrid_bm25_weight=None, + ) + + +def test_chat_continue_cli_with_rag_success(runner, mock_sdk_client): + """Test `owui chat continue` with RAG parameters.""" + mock_updated_chat = MagicMock(spec=models.ChatResponse, id="existing-chat-id", title="Continued RAG Chat") + mock_updated_chat.chat.additional_properties = { + "messages": [ + {"role": "user", "content": "Previous msg."}, + {"role": "assistant", "content": "Previous response."}, + {"role": "user", "content": "New msg about project."}, + {"role": "assistant", "content": "New assist response RAG."}, + ] + } + mock_sdk_client.chats.continue_chat.return_value = mock_updated_chat + + result = runner.invoke( + cli, + [ + "chat", + "continue", + "existing-chat-id", + "What about the new project timeline?", + "--kb-id", "project-kb-1", + "--k", "3" + ] + ) + + assert result.exit_code == 0 + assert "✅ Success! Chat existing-chat-id updated." in result.output + assert "New assist response RAG." in result.output + + mock_sdk_client.chats.continue_chat.assert_awaited_once_with( + chat_id="existing-chat-id", + prompt="What about the new project timeline?", + kb_ids=["project-kb-1"], + k=3, + k_reranker=None, + r=None, + hybrid=None, + hybrid_bm25_weight=None, + ) + +def test_chat_continue_cli_with_no_rag_options(runner, mock_sdk_client): + """Test `owui chat continue` with no RAG options (should pass None for options).""" + mock_updated_chat = MagicMock(spec=models.ChatResponse, id="existing-chat-id", title="No RAG options Chat") + mock_updated_chat.chat.additional_properties = { + "messages": [ + {"role": "user", "content": "Previous msg."}, + {"role": "assistant", "content": "Previous response."}, + {"role": "user", "content": "New msg."}, + {"role": "assistant", "content": "New assist response."}, + ] + } + mock_sdk_client.chats.continue_chat.return_value = mock_updated_chat + + result = runner.invoke( + cli, + [ + "chat", + "continue", + "existing-chat-id", + "Next question." + ] + ) + + assert result.exit_code == 0 + assert "✅ Success! Chat existing-chat-id updated." in result.output + + mock_sdk_client.chats.continue_chat.assert_awaited_once_with( + chat_id="existing-chat-id", + prompt="Next question.", + kb_ids=None, # Should be None if not provided + k=None, + k_reranker=None, + r=None, + hybrid=None, + hybrid_bm25_weight=None, + ) + diff --git a/tests/sdk/test_sdk_chats.py b/tests/sdk/test_sdk_chats.py index f326edd..7351deb 100644 --- a/tests/sdk/test_sdk_chats.py +++ b/tests/sdk/test_sdk_chats.py @@ -1,5 +1,5 @@ import pytest -from unittest.mock import MagicMock +from unittest.mock import MagicMock, AsyncMock from openwebui.exceptions import NotFoundError from openwebui.open_web_ui_client.open_web_ui_client import models @@ -79,8 +79,9 @@ async def test_chats_delete_success(mocker, sdk_client): mock_api_func.assert_awaited_once_with(id="chat-to-delete", client=sdk_client._client) +@pytest.mark.asyncio async def test_chats_create_success(mocker, sdk_client): - """Tests the complex two-step chat creation workflow.""" + """Tests the complex two-step chat creation workflow including optional RAG.""" # 1. Mock the LLM completion API call mock_llm_api = mocker.patch( "openwebui.api.chats.generate_chat_completion_openai_chat_completions_post.asyncio_detailed" @@ -97,7 +98,7 @@ async def test_chats_create_success(mocker, sdk_client): id="new-chat-123", user_id="user1", title="New Chat", - chat={}, + chat=models.ChatResponseChat(), # Chat response often has empty chat on creation created_at=1, updated_at=1, archived=False, @@ -105,99 +106,206 @@ async def test_chats_create_success(mocker, sdk_client): mock_create_response = MagicMock(spec=Response, status_code=200, parsed=final_chat_response) mock_create_api.return_value = mock_create_response - # 3. Call the high-level SDK method - new_chat = await sdk_client.chats.create(model="test-model", prompt="What is the capital of France?") + # --- RAG Specific Mocking --- + # Mock the KnowledgeBaseAPI.query method that ChatsAPI uses + mock_kb_query = mocker.patch( + "openwebui.api.knowledge.KnowledgeBaseAPI.query", new_callable=AsyncMock + ) + # Configure it to return a list of dictionaries, simulating chunks + mock_kb_query.return_value = [ + {"content": "Relevant KB content about France."}, + {"content": "More details about European capitals."}, + ] - # 4. Assertions - assert new_chat.id == "new-chat-123" + # Test 1: Basic chat creation (no RAG) + new_chat_no_rag = await sdk_client.chats.create(model="test-model", prompt="What is the capital of France?") + assert new_chat_no_rag.id == "new-chat-123" + mock_llm_api.assert_awaited_once() # Should be called once for this test + + # Reset mocks for the next test + mock_llm_api.reset_mock() + mock_create_api.reset_mock() + mock_kb_query.reset_mock() + + # Reconfigure mock_create_api for the next chat creation + mock_create_api.return_value = MagicMock(spec=Response, status_code=200, parsed=final_chat_response) + + # Test 2: Chat creation WITH RAG parameters + kb_ids_to_use = ["my-kb-id-1", "another-kb-id-2"] + rag_k = 5 + rag_r = 0.7 + + new_chat_with_rag = await sdk_client.chats.create( + model="test-model", + prompt="Tell me about AI", + kb_ids=kb_ids_to_use, + k=rag_k, + r=rag_r, + hybrid=True, + ) + + assert new_chat_with_rag.id == "new-chat-123" # Still returns the same mocked chat + + # Verify KnowledgeBaseAPI.query was called correctly + mock_kb_query.assert_awaited_once_with( + "Tell me about AI", # Original prompt passed to KB query + kb_ids_to_use, + k=rag_k, + k_reranker=None, # Not set in this call + r=rag_r, + hybrid=True, + hybrid_bm25_weight=None, # Not set in this call + ) - # Assert the LLM call was correct + # Verify LLM was called with the augmented prompt mock_llm_api.assert_awaited_once() llm_call_args = mock_llm_api.call_args[1] - assert llm_call_args["body"]["model"] == "test-model" - assert llm_call_args["body"]["messages"][0]["content"] == "What is the capital of France?" + augmented_user_message = llm_call_args["body"]["messages"][0]["content"] + assert "Please use the following context to answer the question." in augmented_user_message + assert "Relevant KB content about France." in augmented_user_message # Content from mock_kb_query + assert "More details about European capitals." in augmented_user_message + assert "Tell me about AI" in augmented_user_message # Original prompt should be in augmented prompt - # Assert the chat creation call was correct - mock_create_api.assert_awaited_once() + # Verify original prompt (not augmented) was saved in chat history create_call_args = mock_create_api.call_args[1] - - # Verify the payload contains the full conversation sent_messages = create_call_args["body"].chat.additional_properties["messages"] - assert len(sent_messages) == 2 - assert sent_messages[0]["role"] == "user" - assert sent_messages[0]["content"] == "What is the capital of France?" - assert sent_messages[1]["role"] == "assistant" - assert sent_messages[1]["content"] == "The capital of France is Paris." + assert sent_messages[0]["content"] == "Tell me about AI" # Original prompt -# tests/test_sdk_chats.py -@pytest.mark.skip(reason="Broken") +@pytest.mark.asyncio async def test_chats_continue_success(mocker, sdk_client): - """Tests the complex workflow of continuing a chat.""" + """Tests the complex workflow of continuing a chat, including optional RAG.""" chat_id = "existing-chat-123" - # 1. Mock the initial GET request - mock_get_api = mocker.patch("openwebui.api.chats.get_chat_by_id_api_v1_chats_id_get.asyncio_detailed") - initial_chat_data = models.ChatResponse( + # Test 1: Continue chat without RAG + # Define initial_chat_data specifically for this test case + initial_chat_data_no_rag = models.ChatResponse( id=chat_id, user_id="user1", - title="Initial Chat", + title="Initial Chat (No RAG)", chat=models.ChatResponseChat(), created_at=1, updated_at=1, archived=False, ) - initial_chat_data.chat.additional_properties = { + initial_chat_data_no_rag.chat.additional_properties = { "models": ["test-model"], "messages": [{"role": "user", "content": "Hello"}], } - mock_get_response = MagicMock(spec=Response, status_code=200, parsed=initial_chat_data) - mock_get_api.return_value = mock_get_response - # 2. Mock the LLM completion call + mock_get_api = mocker.patch("openwebui.api.chats.get_chat_by_id_api_v1_chats_id_get.asyncio_detailed") + mock_get_api.return_value = MagicMock(spec=Response, status_code=200, parsed=initial_chat_data_no_rag) + mock_llm_api = mocker.patch( - "openwebui.api.chats.generate_chat_completion_openai_chat_completions_post.asyncio_detailed" - ) + "openwebui.api.chats.generate_chat_completion_openai_chat_completions_post.asyncio_detailed") llm_response_data = {"choices": [{"message": {"content": "Hello back!"}}]} - mock_llm_response = MagicMock(spec=Response, status_code=200, parsed=llm_response_data) - mock_llm_api.return_value = mock_llm_response + mock_llm_api.return_value = MagicMock(spec=Response, status_code=200, parsed=llm_response_data) - # 3. Mock the final chat update call - mock_update_api = mocker.patch( - "openwebui.api.chats.update_chat_by_id_api_v1_chats_id_post.asyncio_detailed" + mock_update_api = mocker.patch("openwebui.api.chats.update_chat_by_id_api_v1_chats_id_post.asyncio_detailed") + mock_update_api.return_value = MagicMock(spec=Response, status_code=200, parsed= + models.ChatResponse( + id=chat_id, user_id="user1", title="Updated Chat (No RAG)", chat=models.ChatResponseChat(), created_at=1, + updated_at=2, archived=False ) - final_chat_data = models.ChatResponse( + ) + + mock_kb_query = mocker.patch( # Mock it even if not used, to control behavior + "openwebui.api.knowledge.KnowledgeBaseAPI.query", new_callable=AsyncMock + ) + mock_kb_query.return_value = [] # Ensure no RAG context for this test + + updated_chat_no_rag = await sdk_client.chats.continue_chat(chat_id=chat_id, prompt="How are you?") + assert updated_chat_no_rag.title == "Updated Chat (No RAG)" + mock_get_api.assert_awaited_once_with(id=chat_id, client=sdk_client._client) + mock_llm_api.assert_awaited_once() # Should be called + mock_update_api.assert_awaited_once() # Should be called + mock_kb_query.assert_not_awaited() # Should NOT be called + + # --- Test 2: Continue chat WITH RAG parameters --- + # Reset and reconfigure mocks for this specific test block + mock_get_api.reset_mock() + mock_llm_api.reset_mock() + mock_update_api.reset_mock() + mock_kb_query.reset_mock() + + initial_chat_data_with_rag = models.ChatResponse( id=chat_id, user_id="user1", - title="Updated Chat", + title="Initial Chat (With RAG)", chat=models.ChatResponseChat(), created_at=1, - updated_at=2, + updated_at=1, archived=False, ) - mock_update_response = MagicMock(spec=Response, status_code=200, parsed=final_chat_data) - mock_update_api.return_value = mock_update_response + initial_chat_data_with_rag.chat.additional_properties = { + "models": ["test-model"], + "messages": [{"role": "user", "content": "Previous history message from RAG test."}], + } + mock_get_api.return_value = MagicMock(spec=Response, status_code=200, parsed=initial_chat_data_with_rag) - # Call the SDK method - updated_chat = await sdk_client.chats.continue_chat(chat_id=chat_id, prompt="How are you?") + mock_llm_api.return_value = MagicMock(spec=Response, status_code=200, + parsed=llm_response_data) # Re-use generic LLM response - # Assertions - assert updated_chat.title == "Updated Chat" - mock_get_api.assert_awaited_once_with(id=chat_id, client=sdk_client._client) + mock_update_api.return_value = MagicMock(spec=Response, status_code=200, parsed= + models.ChatResponse( + id=chat_id, user_id="user1", title="Updated Chat (With RAG)", chat=models.ChatResponseChat(), created_at=1, + updated_at=2, archived=False + ) + ) + + # Configure KB query to return content for RAG + mock_kb_query.return_value = [ + {"content": "Relevant KB content for continuation."}, + {"content": "Additional facts for the project."}, + ] - # Verify the payload sent TO the LLM contained the history up to the new prompt + kb_ids_to_use = ["continuation-kb-1"] + rag_k_reranker = 2 + rag_hybrid_bm25_weight = 0.5 + + updated_chat_with_rag = await sdk_client.chats.continue_chat( + chat_id=chat_id, + prompt="Tell me more about the project.", + kb_ids=kb_ids_to_use, + k_reranker=rag_k_reranker, + hybrid_bm25_weight=rag_hybrid_bm25_weight, + ) + + assert updated_chat_with_rag.title == "Updated Chat (With RAG)" # Verify title updated + + # Verify KnowledgeBaseAPI.query was called correctly + mock_kb_query.assert_awaited_once_with( + "Tell me more about the project.", # Original prompt passed to KB query + kb_ids_to_use, + k=None, # Not set in this call + k_reranker=rag_k_reranker, + r=None, # Not set in this call + hybrid=None, # Not set in this call + hybrid_bm25_weight=rag_hybrid_bm25_weight, + ) + + # Verify LLM was called with the augmented prompt including history and new context mock_llm_api.assert_awaited_once() llm_call_args = mock_llm_api.call_args[1] - # FIX: The assertion was slightly incorrect. We check the state of the list - # *when it was passed* to the LLM, which was correctly 2. - assert len(llm_call_args["body"]["messages"]) == 2 - assert llm_call_args["body"]["messages"][0]["content"] == "Hello" - assert llm_call_args["body"]["messages"][1]["content"] == "How are you?" - # Verify the final update call contained the full conversation + # The list of messages passed to LLM should include previous history + the new augmented user message + llm_messages = llm_call_args["body"]["messages"] + assert len(llm_messages) == 2 # "Previous history message from RAG test." + new augmented prompt + + # Check the content of the previous message + assert llm_messages[0]["content"] == "Previous history message from RAG test." + + # Check the content of the augmented user message + augmented_user_message = llm_messages[1]["content"] + assert "Please use the following context to answer the question." in augmented_user_message + assert "Relevant KB content for continuation." in augmented_user_message + assert "Additional facts for the project." in augmented_user_message + assert "Tell me more about the project." in augmented_user_message + + # Verify the final update call saved the original prompt, not the augmented one mock_update_api.assert_awaited_once() update_call_args = mock_update_api.call_args[1] final_messages = update_call_args["body"].chat.additional_properties["messages"] - assert len(final_messages) == 3 # The user prompt, new user prompt, and assistant response - assert final_messages[2]["role"] == "assistant" - assert final_messages[2]["content"] == "Hello back!" + assert len(final_messages) == 2 # Original chat and new message + assert final_messages[0]["content"] == "Previous history message from RAG test." + assert final_messages[1]["content"] == "Tell me more about the project." # Should be original prompt \ No newline at end of file diff --git a/tests/sdk/test_sdk_knowledge.py b/tests/sdk/test_sdk_knowledge.py index 701cfc8..e074059 100644 --- a/tests/sdk/test_sdk_knowledge.py +++ b/tests/sdk/test_sdk_knowledge.py @@ -1,8 +1,9 @@ +import httpx import pytest from unittest.mock import MagicMock, AsyncMock from openwebui.exceptions import APIError from openwebui.open_web_ui_client.open_web_ui_client import models # Import generated models -from openwebui.open_web_ui_client.open_web_ui_client.types import Response +from openwebui.open_web_ui_client.open_web_ui_client.types import Response, UNSET import json # Import json for json.dumps here from pathlib import Path # Import Path from pathlib here @@ -343,3 +344,152 @@ async def test_knowledge_list_files_success(mocker, sdk_client): assert files[0].id == "file1" assert files[1].meta["name"] == "report.docx" # FIX: Access meta as a dict directly mock_list_files_api_call.assert_awaited_once_with(id="test-kb-id", client=sdk_client._client) + + +@pytest.mark.asyncio +async def test_knowledge_query_success_basic(mocker, sdk_client): + """Test successful basic query to a single knowledge base.""" + mock_query_api_call = mocker.patch( + "openwebui.api.knowledge.query_kb_api_call", new_callable=AsyncMock + ) + + # Simulate a successful retrieval response with content + mock_retrieved_chunks = [ + {"content": "This is a test document snippet.", "meta": {"file": "test.txt"}}, + {"content": "Another relevant piece of information.", "meta": {"file": "another.pdf"}}, + ] + mock_response_object = MagicMock(spec=Response, status_code=200, parsed=mock_retrieved_chunks) + mock_query_api_call.return_value = mock_response_object + + query_text = "What is the main topic?" + kb_ids = ["kb-id-1"] + k_value = 2 + + result_chunks = await sdk_client.knowledge.query(query_text, kb_ids, k=k_value) + + # Assert that the low-level API was called with the correct QueryCollectionsForm + mock_query_api_call.assert_awaited_once() + call_args, call_kwargs = mock_query_api_call.call_args + assert isinstance(call_kwargs['body'], models.QueryCollectionsForm) + assert call_kwargs['body'].collection_names == kb_ids + assert call_kwargs['body'].query == query_text + assert call_kwargs['body'].k == k_value + assert call_kwargs['body'].k_reranker is UNSET + assert call_kwargs['body'].r is UNSET + assert call_kwargs['body'].hybrid is UNSET + assert call_kwargs['body'].hybrid_bm25_weight is UNSET + + # Assert the returned data matches the mock response + assert isinstance(result_chunks, list) + assert len(result_chunks) == len(mock_retrieved_chunks) + assert result_chunks[0]["content"] == "This is a test document snippet." + assert result_chunks[1]["meta"]["file"] == "another.pdf" + + +@pytest.mark.asyncio +async def test_knowledge_query_success_multiple_kbs_and_all_rag_params(mocker, sdk_client): + """Test successful query to multiple knowledge bases with all RAG parameters.""" + mock_query_api_call = mocker.patch( + "openwebui.api.knowledge.query_kb_api_call", new_callable=AsyncMock + ) + + mock_retrieved_chunks = [ + {"content": "Content from KB A.", "meta": {"kb": "KB_A"}}, + {"content": "Content from KB B.", "meta": {"kb": "KB_B"}}, + ] + mock_response_object = MagicMock(spec=Response, status_code=200, parsed=mock_retrieved_chunks) + mock_query_api_call.return_value = mock_response_object + + query_text = "Advanced RAG search" + kb_ids = ["kb-alpha", "kb-beta"] + k_value = 5 + k_reranker_value = 3 + r_value = 0.75 + hybrid_value = True + hybrid_bm25_weight_value = 0.3 + + result_chunks = await sdk_client.knowledge.query( + query_text, + kb_ids, + k=k_value, + k_reranker=k_reranker_value, + r=r_value, + hybrid=hybrid_value, + hybrid_bm25_weight=hybrid_bm25_weight_value + ) + + # Assert the low-level API was called with the correct QueryCollectionsForm + mock_query_api_call.assert_awaited_once() + call_args, call_kwargs = mock_query_api_call.call_args + assert isinstance(call_kwargs['body'], models.QueryCollectionsForm) + assert call_kwargs['body'].collection_names == kb_ids + assert call_kwargs['body'].query == query_text + assert call_kwargs['body'].k == k_value + assert call_kwargs['body'].k_reranker == k_reranker_value + assert call_kwargs['body'].r == r_value + assert call_kwargs['body'].hybrid == hybrid_value + assert call_kwargs['body'].hybrid_bm25_weight == hybrid_bm25_weight_value + + assert isinstance(result_chunks, list) + assert len(result_chunks) == len(mock_retrieved_chunks) + + +@pytest.mark.asyncio +async def test_knowledge_query_empty_result(mocker, sdk_client): + """Test query returning an empty list of chunks.""" + mock_query_api_call = mocker.patch( + "openwebui.api.knowledge.query_kb_api_call", new_callable=AsyncMock + ) + # Simulate an empty retrieval result + mock_response_object = MagicMock(spec=Response, status_code=200, parsed=[]) + mock_query_api_call.return_value = mock_response_object + + query_text = "Non-existent topic" + kb_ids = ["empty-kb"] + + result_chunks = await sdk_client.knowledge.query(query_text, kb_ids) + + mock_query_api_call.assert_awaited_once() + assert result_chunks == [] + + +@pytest.mark.asyncio +async def test_knowledge_query_api_error(mocker, sdk_client): + """Test APIError during knowledge base query.""" + mock_query_api_call = mocker.patch( + "openwebui.api.knowledge.query_kb_api_call", new_callable=AsyncMock + ) + + # Simulate an API error response (e.g., 400 Bad Request) + mock_bad_response = MagicMock(spec=Response, status_code=400, content=b'{"detail": "Invalid query"}') + mock_bad_response.parsed = None # Ensure handle_api_response goes to raw content + mock_query_api_call.return_value = mock_bad_response + + query_text = "Error query" + kb_ids = ["invalid-kb"] + + with pytest.raises(APIError) as exc_info: + await sdk_client.knowledge.query(query_text, kb_ids) + + assert exc_info.value.status_code == 400 + assert "Invalid query" in str(exc_info.value) + mock_query_api_call.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_knowledge_query_connection_error(mocker, sdk_client): + """Test ConnectionError during knowledge base API call.""" + mock_query_api_call = mocker.patch( + "openwebui.api.knowledge.query_kb_api_call", new_callable=AsyncMock + ) + # Simulate a network connection error + mock_query_api_call.side_effect = httpx.ConnectError("Connection refused") + + query_text = "Network issue" + kb_ids = ["remote-kb"] + + with pytest.raises(ConnectionError) as exc_info: + await sdk_client.knowledge.query(query_text, kb_ids) + + assert "Connection refused" in str(exc_info.value) + mock_query_api_call.assert_awaited_once() \ No newline at end of file From d9e1579231d1812f15626afc7490833f78b76848 Mon Sep 17 00:00:00 2001 From: Herman Haggerty Date: Mon, 1 Sep 2025 15:59:02 -0500 Subject: [PATCH 2/3] Update openwebui/cli/main.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- openwebui/cli/main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/openwebui/cli/main.py b/openwebui/cli/main.py index 4279912..709884d 100644 --- a/openwebui/cli/main.py +++ b/openwebui/cli/main.py @@ -23,7 +23,6 @@ handler.setFormatter(formatter) log.addHandler(handler) -logging.basicConfig(level=logging.DEBUG) # --- Output Formatting Helper --- def format_output(data: Any, output_format: str): From 5515465eabede081e62feeb4de54a5630c8ecf3d Mon Sep 17 00:00:00 2001 From: Herman Haggerty Date: Mon, 1 Sep 2025 15:59:25 -0500 Subject: [PATCH 3/3] Update openwebui/api/chats.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- openwebui/api/chats.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openwebui/api/chats.py b/openwebui/api/chats.py index 7f85219..dbaf5ec 100644 --- a/openwebui/api/chats.py +++ b/openwebui/api/chats.py @@ -56,7 +56,7 @@ async def create( # --- RAG WORKFLOW ADDITION --- if kb_ids: - print(kb_ids) + log.debug(f"kb_ids provided for RAG workflow: {kb_ids}") log.info(f"RAG workflow activated. Querying KBs {kb_ids} for context.") try: retrieved_chunks = await self.knowledge.query(