This commit is contained in:
Mark Campbell 2025-06-27 11:39:51 +02:00 committed by GitHub
commit 36b44655c6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 52 additions and 3 deletions

View file

@ -14676,7 +14676,8 @@
"description": "Template for formatting each retrieved chunk in the context. Available placeholders: {index} (1-based chunk ordinal), {chunk.content} (chunk content string), {metadata} (chunk metadata dict). Default: \"Result {index}\\nContent: {chunk.content}\\nMetadata: {metadata}\\n\"" "description": "Template for formatting each retrieved chunk in the context. Available placeholders: {index} (1-based chunk ordinal), {chunk.content} (chunk content string), {metadata} (chunk metadata dict). Default: \"Result {index}\\nContent: {chunk.content}\\nMetadata: {metadata}\\n\""
}, },
"mode": { "mode": {
"type": "string", "$ref": "#/components/schemas/RAGSearchMode",
"default": "vector",
"description": "Search mode for retrieval—either \"vector\", \"keyword\", or \"hybrid\". Default \"vector\"." "description": "Search mode for retrieval—either \"vector\", \"keyword\", or \"hybrid\". Default \"vector\"."
}, },
"ranker": { "ranker": {
@ -14711,6 +14712,16 @@
} }
} }
}, },
"RAGSearchMode": {
"type": "string",
"enum": [
"vector",
"keyword",
"hybrid"
],
"title": "RAGSearchMode",
"description": "Search modes for RAG query retrieval: - VECTOR: Uses vector similarity search for semantic matching - KEYWORD: Uses keyword-based search for exact matching - HYBRID: Combines both vector and keyword search for better results"
},
"RRFRanker": { "RRFRanker": {
"type": "object", "type": "object",
"properties": { "properties": {

View file

@ -10260,7 +10260,8 @@ components:
content string), {metadata} (chunk metadata dict). Default: "Result {index}\nContent: content string), {metadata} (chunk metadata dict). Default: "Result {index}\nContent:
{chunk.content}\nMetadata: {metadata}\n" {chunk.content}\nMetadata: {metadata}\n"
mode: mode:
type: string $ref: '#/components/schemas/RAGSearchMode'
default: vector
description: >- description: >-
Search mode for retrieval—either "vector", "keyword", or "hybrid". Default Search mode for retrieval—either "vector", "keyword", or "hybrid". Default
"vector". "vector".
@ -10287,6 +10288,17 @@ components:
mapping: mapping:
default: '#/components/schemas/DefaultRAGQueryGeneratorConfig' default: '#/components/schemas/DefaultRAGQueryGeneratorConfig'
llm: '#/components/schemas/LLMRAGQueryGeneratorConfig' llm: '#/components/schemas/LLMRAGQueryGeneratorConfig'
RAGSearchMode:
type: string
enum:
- vector
- keyword
- hybrid
title: RAGSearchMode
description: >-
Search modes for RAG query retrieval: - VECTOR: Uses vector similarity search
for semantic matching - KEYWORD: Uses keyword-based search for exact matching
- HYBRID: Combines both vector and keyword search for better results
RRFRanker: RRFRanker:
type: object type: object
properties: properties:

View file

@ -87,6 +87,20 @@ class RAGQueryGenerator(Enum):
custom = "custom" custom = "custom"
@json_schema_type
class RAGSearchMode(Enum):
"""
Search modes for RAG query retrieval:
- VECTOR: Uses vector similarity search for semantic matching
- KEYWORD: Uses keyword-based search for exact matching
- HYBRID: Combines both vector and keyword search for better results
"""
VECTOR = "vector"
KEYWORD = "keyword"
HYBRID = "hybrid"
@json_schema_type @json_schema_type
class DefaultRAGQueryGeneratorConfig(BaseModel): class DefaultRAGQueryGeneratorConfig(BaseModel):
type: Literal["default"] = "default" type: Literal["default"] = "default"
@ -128,7 +142,7 @@ class RAGQueryConfig(BaseModel):
max_tokens_in_context: int = 4096 max_tokens_in_context: int = 4096
max_chunks: int = 5 max_chunks: int = 5
chunk_template: str = "Result {index}\nContent: {chunk.content}\nMetadata: {metadata}\n" chunk_template: str = "Result {index}\nContent: {chunk.content}\nMetadata: {metadata}\n"
mode: str | None = None mode: RAGSearchMode | None = RAGSearchMode.VECTOR
ranker: Ranker | None = Field(default=None) # Only used for hybrid mode ranker: Ranker | None = Field(default=None) # Only used for hybrid mode
@field_validator("chunk_template") @field_validator("chunk_template")

View file

@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest import pytest
from llama_stack.apis.tools.rag_tool import RAGQueryConfig
from llama_stack.apis.vector_io import ( from llama_stack.apis.vector_io import (
Chunk, Chunk,
ChunkMetadata, ChunkMetadata,
@ -60,3 +61,14 @@ class TestRagQuery:
) )
assert expected_metadata_string in result.content[1].text assert expected_metadata_string in result.content[1].text
assert result.content is not None assert result.content is not None
async def test_query_raises_incorrect_mode(self):
with pytest.raises(ValueError):
RAGQueryConfig(mode="invalid_mode")
@pytest.mark.asyncio
async def test_query_accepts_valid_modes(self):
RAGQueryConfig() # Test default (vector)
RAGQueryConfig(mode="vector") # Test vector
RAGQueryConfig(mode="keyword") # Test keyword
RAGQueryConfig(mode="hybrid") # Test hybrid