Implement Chat with RAG

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

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.

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
@@ -24,7 +24,6 @@
     log.addHandler(handler)
 
 
-
 # --- Output Formatting Helper ---
 def format_output(data: Any, output_format: str):
     """
@@ -97,36 +96,93 @@ 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"])
         else:
             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 +195,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 +286,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