MAIT-205: Replace ROAT Filter Function with OWUI Workspace Tool #48
+237
−34
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,203 @@ | ||
| """ | ||
| title: ROAT Knowledge Base Search | ||
|
Member
There was a problem hiding this comment. Please rename this to just the "Knowledge Base Search" |
||
| author: WikiTeq | ||
| date: 2025-05-01 | ||
| version: 1.0 | ||
| license: MIT | ||
| description: Searches the RAG-of-All-Trades knowledge base and returns relevant context for the user's query. | ||
| requirements: requests | ||
| """ | ||
|
|
||
| import asyncio | ||
| import logging | ||
| import re | ||
| from collections.abc import Awaitable, Callable | ||
| from typing import Optional | ||
|
|
||
| import requests | ||
| from pydantic import BaseModel, Field | ||
|
|
||
| log = logging.getLogger(__name__) | ||
| log.setLevel(logging.INFO) | ||
|
|
||
|
|
||
| def _parse_raw_chunk(raw_text: str) -> dict: | ||
| match = re.match(r"Score:\s*([\d.]+)\s*\|\s*Text:\s*(.*)", raw_text, re.DOTALL) | ||
| if match: | ||
| return {"score": float(match.group(1)), "text": match.group(2).strip()} | ||
| return {"score": 0.0, "text": raw_text.strip()} | ||
|
|
||
|
|
||
| def _get_filename_from_extras(extras: dict) -> Optional[str]: | ||
| return extras.get("key") or extras.get("filename") or extras.get("name") or None | ||
|
|
||
|
|
||
| def _call_rag_service(url: str, api_key: str, timeout: int, top_k: int, query: str) -> dict: | ||
| payload = {"query": query, "top_k": top_k, "metadata_filters": {}} | ||
| headers = {"Content-Type": "application/json"} | ||
| if api_key: | ||
| headers["Authorization"] = f"Bearer {api_key}" | ||
| url = url.strip() | ||
| log.info("Calling ROAT: query_length=%d", len(query)) | ||
| response = requests.post(url, json=payload, headers=headers, timeout=timeout) | ||
| response.raise_for_status() | ||
| return response.json() | ||
|
|
||
|
|
||
| def _format_context_and_sources(rag_result: dict) -> tuple[str, list]: | ||
| references = rag_result.get("references", []) or [] | ||
| raw_chunks = rag_result.get("raw") or [] | ||
|
|
||
| if not references and not raw_chunks: | ||
| return "", [] | ||
|
|
||
| context_parts = [] | ||
| sources = [] | ||
| _internal_fields = {"key", "format", "version", "checksum"} | ||
|
|
||
| for i in range(max(len(references), len(raw_chunks))): | ||
| ref = references[i] if i < len(references) else {} | ||
| extras = ref.get("extras") or {} | ||
| score = ref.get("score", 0.0) | ||
|
|
||
| text = ref.get("text", "") | ||
| if not text and i < len(raw_chunks): | ||
| parsed = _parse_raw_chunk(raw_chunks[i]) | ||
| text = parsed["text"] | ||
| if score == 0.0: | ||
| score = parsed["score"] | ||
|
|
||
| if not text: | ||
| continue | ||
|
|
||
| filename = _get_filename_from_extras(extras) | ||
| source_name = ( | ||
| ref.get("title") | ||
| or ref.get("source_name") | ||
| or filename | ||
| or f"Source {i + 1}" | ||
| ) | ||
|
|
||
| metadata_fields = {k: v for k, v in extras.items() if k not in _internal_fields} | ||
| metadata_fields["url"] = ref.get("url") or extras.get("url") | ||
| metadata_md = "\n".join( | ||
| f"- *{k}*: {v}" for k, v in metadata_fields.items() if v is not None | ||
| ) | ||
| metadata_section = f"## Metadata\n\n{metadata_md}" if metadata_md else "" | ||
|
|
||
| context_parts.append(f"[Source: {source_name}]\n\n{metadata_section}\n\n{text}\n") | ||
|
|
||
| source_obj = { | ||
| "source": {"name": source_name}, | ||
| "document": [text[:1000] if len(text) > 1000 else text], | ||
|
Member
There was a problem hiding this comment. While we're at it, can we please make this configurable via Valve, making it unlimited by default? |
||
| "metadata": [ | ||
| { | ||
| "source": source_name, | ||
| "file": filename, | ||
| "relevance_score": score, | ||
|
Member
There was a problem hiding this comment. Not something new - it was like that in the Filter function, but I noticed that the relevance_score is not displayed by the OWUI when viewing the citation source. Do you have any ideas why? |
||
| "type": extras.get("format", "document"), | ||
| "storage": extras.get("source"), | ||
| "key": extras.get("key"), | ||
| "checksum": extras.get("checksum"), | ||
| "version": extras.get("version"), | ||
| "format": extras.get("format"), | ||
| } | ||
| ], | ||
| } | ||
| url = ref.get("url") or extras.get("url") | ||
| if url: | ||
| source_obj["source"]["url"] = url | ||
|
|
||
| sources.append(source_obj) | ||
|
|
||
| return "\n".join(context_parts), sources | ||
|
|
||
|
|
||
| class Tools: | ||
| class Valves(BaseModel): | ||
| rag_service_url: str = Field( | ||
| default="", | ||
| description="Full URL to the ROAT query endpoint, e.g. http://api:8000/query.", | ||
| ) | ||
| rag_service_api_key: str = Field( | ||
| default="", | ||
| description="Bearer token for the ROAT API (leave blank if not required).", | ||
| ) | ||
| rag_service_timeout: int = Field( | ||
| default=30, | ||
| description="Request timeout in seconds.", | ||
| ) | ||
| top_k: int = Field( | ||
| default=20, | ||
| description="Number of top results to retrieve from the knowledge base.", | ||
| ) | ||
|
|
||
| def __init__(self): | ||
| self.valves = self.Valves() | ||
|
|
||
| async def search_knowledge_base( | ||
| self, | ||
| query: str, | ||
| __event_emitter__: Callable[[dict], Awaitable[None]] | None = None, | ||
| ) -> str: | ||
| """ | ||
| Search the mAItion knowledge base for information relevant to the user's question. | ||
|
|
||
| Use this tool when: | ||
| - The user asks a factual question that likely requires domain-specific knowledge | ||
| - The user asks about documented processes, policies, data, or internal content | ||
| - The user's question cannot be confidently answered from general training knowledge alone | ||
|
Member
There was a problem hiding this comment. In fact, we want this tool used for most of the queries, and we want to disincentivize the LLM from using general knowledge. Thus, I propose to drop this line and instead replace it with something that would instruct the model to use this tool more frequently We're of course just guessing for now until we have some evaluations in place |
||
| - The user explicitly asks to search the knowledge base or documentation | ||
|
|
||
| Do NOT use this tool for: | ||
| - Casual greetings or small talk | ||
| - Simple arithmetic or general world knowledge | ||
| - Follow-up questions within an ongoing conversation where context was already retrieved | ||
|
Member
There was a problem hiding this comment. I don't think this specific instruction as follow up questions may require queries into the ROAT, let's try to reformulate this |
||
|
|
||
| Args: | ||
| query: A concise search query derived from the user's question | ||
|
|
||
| Returns: | ||
| Retrieved context passages from the knowledge base, or a message that nothing was found. | ||
|
Member
There was a problem hiding this comment. Please reword into |
||
| """ | ||
|
|
||
| async def emit(description: str, done: bool = False) -> None: | ||
| if __event_emitter__: | ||
| await __event_emitter__( | ||
| {"type": "status", "data": {"description": description, "done": done}} | ||
| ) | ||
|
|
||
| if not self.valves.rag_service_url: | ||
| await emit("Knowledge base URL is not configured in Tool Valves.", done=True) | ||
| return "Error: rag_service_url is not configured in Tool Valves." | ||
|
|
||
| await emit("Searching knowledge base…") | ||
|
|
||
| try: | ||
| log.info("ROAT url=%r top_k=%r timeout=%r", self.valves.rag_service_url, self.valves.top_k, self.valves.rag_service_timeout) | ||
| rag_result = await asyncio.to_thread( | ||
| _call_rag_service, | ||
| self.valves.rag_service_url, | ||
| self.valves.rag_service_api_key, | ||
| self.valves.rag_service_timeout, | ||
| self.valves.top_k, | ||
| query, | ||
| ) | ||
|
coderabbitai[bot] marked this conversation as resolved.
|
||
| except Exception as e: | ||
| log.error("ROAT request failed: %s", e, exc_info=True) | ||
| await emit("Failed to reach the knowledge base.", done=True) | ||
| return "Error: could not reach the knowledge base. Check the server logs for details." | ||
|
|
||
| context, sources = _format_context_and_sources(rag_result) | ||
|
|
||
| if not context: | ||
| await emit("No relevant information found.", done=True) | ||
| return "No relevant information was found in the knowledge base for this query." | ||
|
|
||
| if __event_emitter__: | ||
| for src in sources: | ||
| await __event_emitter__({"type": "source", "data": src}) | ||
|
|
||
| await emit(f"Found {len(sources)} relevant source(s).", done=True) | ||
| log.info("Returning context with %d sources (%d chars)", len(sources), len(context)) | ||
| return context | ||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this fits more to be included in a changelog (what we don't yet have) rather than in a README. Let's instead simply ensure both the Filter function and the Tool are documented