Merge pull request #868 from Anxhela21/anx/streaming_query

LCORE-873: Support Solr Vector I/O Provider in LCORE

Author: tisnik

Makefile

@@ -33,7 +33,7 @@ test-integration: ## Run integration tests tests
 	COVERAGE_FILE="${ARTIFACT_DIR}/.coverage.integration" uv run python -m pytest tests/integration --cov=src --cov-report term-missing --cov-report "json:${ARTIFACT_DIR}/coverage_integration.json" --junit-xml="${ARTIFACT_DIR}/junit_integration.xml" --cov-fail-under=10
 
 test-e2e: ## Run end to end tests for the service
-	script -q -e -c "uv run behave --color --format pretty --tags=-skip -D dump_errors=true @tests/e2e/test_list.txt"
+	uv run behave --color --format pretty --tags=-skip -D dump_errors=true @tests/e2e/test_list.txt
 
 test-e2e-local: ## Run end to end tests for the service
 	uv run behave --color --format pretty --tags=-skip -D dump_errors=true @tests/e2e/test_list.txt

lightspeed-stack.yaml

@@ -30,3 +30,8 @@ conversation_cache:
 
 authentication:
   module: "noop"
+
+
+solr:
+  enabled: true
+  offline: true

run.yaml

@@ -15,7 +15,7 @@ apis:
 benchmarks: []
 datasets: []
 image_name: starter
-# external_providers_dir: /opt/app-root/src/.llama/providers.d
+external_providers_dir: ${env.EXTERNAL_PROVIDERS_DIR}
 
 providers:
   inference:
@@ -24,7 +24,9 @@ providers:
     config:
       api_key: ${env.OPENAI_API_KEY}
       allowed_models: ["${env.E2E_OPENAI_MODEL:=gpt-4o-mini}"]
-  - config: {}
+  - config:
+      allowed_models:
+        - ${env.EMBEDDING_MODEL_DIR}
     provider_id: sentence-transformers
     provider_type: inline::sentence-transformers
   files:
@@ -56,6 +58,18 @@ providers:
     provider_id: rag-runtime
     provider_type: inline::rag-runtime
   vector_io:
+  - provider_id: solr-vector
+    provider_type: remote::solr_vector_io
+    config:
+      solr_url: http://localhost:8983/solr
+      collection_name: portal-rag
+      vector_field: chunk_vector
+      content_field: chunk
+      embedding_dimension: 384
+      embedding_model: ${env.EMBEDDING_MODEL_DIR}
+      persistence:
+        namespace: portal-rag
+        backend: kv_default
   - config: # Define the storage backend for RAG
       persistence:
         namespace: vector_io::faiss
@@ -127,7 +141,13 @@ storage:
       namespace: prompts
       backend: kv_default
 registered_resources:
-  models: []
+  models:
+  - model_id: granite-embedding-30m
+    model_type: embedding
+    provider_id: sentence-transformers
+    provider_model_id: ${env.EMBEDDING_MODEL_DIR}
+    metadata:
+      embedding_dimension: 384
   shields:
   - shield_id: llama-guard
     provider_id: llama-guard

src/app/endpoints/a2a.py

@@ -311,6 +311,7 @@ async def _process_task_streaming(  # pylint: disable=too-many-locals
             generate_topic_summary=True,
             media_type=None,
             vector_store_ids=vector_store_ids,
+            solr=None,
         )
 
         # Get LLM client and select model

src/app/endpoints/query.py

@@ -2,8 +2,8 @@
 
 """Handler for REST API call to provide answer to query using Response API."""
 
+import datetime
 import logging
-from datetime import UTC, datetime
 from typing import Annotated, Any, cast
 
 from fastapi import APIRouter, Depends, HTTPException, Request
@@ -25,13 +25,15 @@
 from configuration import configuration
 from models.config import Action
 from models.requests import QueryRequest
+
 from models.responses import (
     ForbiddenResponse,
     InternalServerErrorResponse,
     NotFoundResponse,
     PromptTooLongResponse,
     QueryResponse,
     QuotaExceededResponse,
+    ReferencedDocument,
     ServiceUnavailableResponse,
     UnauthorizedResponse,
     UnprocessableEntityResponse,
@@ -63,7 +65,11 @@
     run_shield_moderation,
 )
 from utils.suid import normalize_conversation_id
-from utils.types import ResponsesApiParams, TurnSummary
+from utils.types import (
+    ResponsesApiParams,
+    TurnSummary,
+)
+from utils.vector_search import perform_vector_search, format_rag_context_for_injection
 
 logger = logging.getLogger("app.endpoints.handlers")
 router = APIRouter(tags=["query"])
@@ -77,9 +83,9 @@
         examples=["endpoint", "conversation read", "model override"]
     ),
     404: NotFoundResponse.openapi_response(
-        examples=["model", "conversation", "provider"]
+        examples=["conversation", "model", "provider"]
     ),
-    413: PromptTooLongResponse.openapi_response(),
+    # 413: PromptTooLongResponse.openapi_response(),
     422: UnprocessableEntityResponse.openapi_response(),
     429: QuotaExceededResponse.openapi_response(),
     500: InternalServerErrorResponse.openapi_response(examples=["configuration"]),
@@ -117,7 +123,7 @@ async def query_endpoint_handler(
     """
     check_configuration_loaded(configuration)
 
-    started_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+    started_at = datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
     user_id, _, _skip_userid_check, token = auth
     # Check token availability
     check_tokens_available(configuration.quota_limiters, user_id)
@@ -145,6 +151,19 @@ async def query_endpoint_handler(
 
     client = AsyncLlamaStackClientHolder().get_client()
 
+    doc_ids_from_chunks: list[ReferencedDocument] = []
+    pre_rag_chunks: list[Any] = []  # use your RAGChunk type (or the upstream one)
+
+    _, _, doc_ids_from_chunks, pre_rag_chunks = await perform_vector_search(
+        client, query_request, configuration
+    )
+
+    rag_context = format_rag_context_for_injection(pre_rag_chunks)
+    if rag_context:
+        # safest: mutate a local copy so we don't surprise other logic
+        query_request = query_request.model_copy(deep=True)  # pydantic v2
+        query_request.query = query_request.query + rag_context
+
     # Prepare API request parameters
     responses_params = await prepare_responses_params(
         client,
@@ -168,6 +187,14 @@ async def query_endpoint_handler(
     # Retrieve response using Responses API
     turn_summary = await retrieve_response(client, responses_params)
 
+    if pre_rag_chunks:
+        turn_summary.rag_chunks = pre_rag_chunks + (turn_summary.rag_chunks or [])
+
+    if doc_ids_from_chunks:
+        turn_summary.referenced_documents = parse_referenced_docs(
+            doc_ids_from_chunks + (turn_summary.referenced_documents or [])
+        )
+
     # Get topic summary for new conversation
     if not user_conversation and query_request.generate_topic_summary:
         logger.debug("Generating topic summary for new conversation")
@@ -190,7 +217,9 @@ async def query_endpoint_handler(
         quota_limiters=configuration.quota_limiters, user_id=user_id
     )
 
-    completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+    completed_at = datetime.datetime.now(datetime.timezone.utc).strftime(
+        "%Y-%m-%dT%H:%M:%SZ"
+    )
     conversation_id = normalize_conversation_id(responses_params.conversation)
 
     logger.info("Storing query results")
@@ -222,6 +251,21 @@ async def query_endpoint_handler(
     )
 
 
+def parse_referenced_docs(
+    docs: list[ReferencedDocument],
+) -> list[ReferencedDocument]:
+    """Remove duplicate referenced documents based on URL and title."""
+    seen: set[tuple[str | None, str | None]] = set()
+    out: list[ReferencedDocument] = []
+    for d in docs:
+        key = (str(d.doc_url) if d.doc_url else None, d.doc_title)
+        if key in seen:
+            continue
+        seen.add(key)
+        out.append(d)
+    return out
+
+
 async def retrieve_response(  # pylint: disable=too-many-locals
     client: AsyncLlamaStackClient,
     responses_params: ResponsesApiParams,

src/app/endpoints/streaming_query.py

@@ -1,8 +1,8 @@
 """Streaming query handler using Responses API."""
 
+import datetime
 import json
 import logging
-from datetime import UTC, datetime
 from typing import Annotated, Any, AsyncIterator, Optional, cast
 
 from fastapi import APIRouter, Depends, HTTPException, Request
@@ -53,7 +53,7 @@
     UnauthorizedResponse,
     UnprocessableEntityResponse,
 )
-from utils.types import ReferencedDocument
+from utils.types import RAGChunk, ReferencedDocument
 from utils.endpoints import (
     check_configuration_loaded,
     validate_and_retrieve_conversation,
@@ -85,6 +85,7 @@
 from utils.suid import normalize_conversation_id
 from utils.token_counter import TokenCounter
 from utils.types import ResponsesApiParams, TurnSummary
+from utils.vector_search import format_rag_context_for_injection, perform_vector_search
 
 logger = logging.getLogger(__name__)
 router = APIRouter(tags=["streaming_query"])
@@ -100,7 +101,7 @@
     404: NotFoundResponse.openapi_response(
         examples=["conversation", "model", "provider"]
     ),
-    413: PromptTooLongResponse.openapi_response(),
+    # 413: PromptTooLongResponse.openapi_response(),
     422: UnprocessableEntityResponse.openapi_response(),
     429: QuotaExceededResponse.openapi_response(),
     500: InternalServerErrorResponse.openapi_response(examples=["configuration"]),
@@ -144,7 +145,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
     check_configuration_loaded(configuration)
 
     user_id, _user_name, _skip_userid_check, token = auth
-    started_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+    started_at = datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
 
     # Check token availability
     check_tokens_available(configuration.quota_limiters, user_id)
@@ -172,6 +173,17 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
 
     client = AsyncLlamaStackClientHolder().get_client()
 
+    pre_rag_chunks: list[RAGChunk] = []
+    doc_ids_from_chunks: list[ReferencedDocument] = []
+
+    _, _, doc_ids_from_chunks, pre_rag_chunks = await perform_vector_search(
+        client, query_request, configuration
+    )
+    rag_context = format_rag_context_for_injection(pre_rag_chunks)
+    if rag_context:
+        query_request = query_request.model_copy(deep=True)
+        query_request.query = query_request.query + rag_context
+
     # Prepare API request parameters
     responses_params = await prepare_responses_params(
         client=client,
@@ -212,6 +224,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
     generator, turn_summary = await retrieve_response_generator(
         responses_params=responses_params,
         context=context,
+        doc_ids_from_chunks=doc_ids_from_chunks,
     )
 
     response_media_type = (
@@ -234,6 +247,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
 async def retrieve_response_generator(
     responses_params: ResponsesApiParams,
     context: ResponseGeneratorContext,
+    doc_ids_from_chunks: list[ReferencedDocument],
 ) -> tuple[AsyncIterator[str], TurnSummary]:
     """
     Retrieve the appropriate response generator.
@@ -273,6 +287,8 @@ async def retrieve_response_generator(
         response = await context.client.responses.create(
             **responses_params.model_dump()
         )
+        # Store pre-RAG documents for later merging
+        turn_summary.pre_rag_documents = doc_ids_from_chunks
         return response_generator(response, context, turn_summary), turn_summary
 
     # Handle know LLS client errors only at stream creation time and shield execution
@@ -373,7 +389,7 @@ async def generate_response(
         turn_summary.referenced_documents,
         context.query_request.media_type or MEDIA_TYPE_JSON,
     )
-    completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+    completed_at = datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
 
     # Store query results (transcript, conversation details, cache)
     logger.info("Storing query results")
@@ -571,9 +587,21 @@ async def response_generator(  # pylint: disable=too-many-branches,too-many-stat
     turn_summary.token_usage = extract_token_usage(
         latest_response_object, context.model_id
     )
-    turn_summary.referenced_documents = parse_referenced_documents(
-        latest_response_object
-    )
+    tool_based_documents = parse_referenced_documents(latest_response_object)
+
+    # Merge pre-RAG documents with tool-based documents (similar to query.py)
+    if turn_summary.pre_rag_documents:
+        all_documents = turn_summary.pre_rag_documents + tool_based_documents
+        seen = set()
+        deduplicated_documents = []
+        for doc in all_documents:
+            key = (doc.doc_url, doc.doc_title)
+            if key not in seen:
+                seen.add(key)
+                deduplicated_documents.append(doc)
+        turn_summary.referenced_documents = deduplicated_documents
+    else:
+        turn_summary.referenced_documents = tool_based_documents
 
 
 def stream_http_error_event(
@@ -608,7 +636,7 @@ def stream_http_error_event(
 
 def format_stream_data(d: dict) -> str:
     """
-    Format a dictionary as a Server-Sent Events (SSE) data string.
+    Create a response generator function for Responses API streaming.
 
     Parameters:
         d (dict): The data to be formatted as an SSE event.
@@ -694,22 +722,24 @@ def stream_event(data: dict, event_type: str, media_type: str) -> str:
     """Build an item to yield based on media type.
 
     Args:
-        data: The data to yield.
-        event_type: The type of event (e.g. token, tool request, tool execution).
-        media_type: Media type of the response (e.g. text or JSON).
+        data: Dictionary containing the event data
+        event_type: Type of event (token, tool call, etc.)
+        media_type: The media type for the response format
 
     Returns:
-        str: The formatted string or JSON to yield.
+        SSE-formatted string representing the event
     """
     if media_type == MEDIA_TYPE_TEXT:
         if event_type == LLM_TOKEN_EVENT:
-            return data["token"]
+            return data.get("token", "")
         if event_type == LLM_TOOL_CALL_EVENT:
-            return f"\nTool call: {json.dumps(data)}\n"
+            return f"[Tool Call: {data.get('function_name', 'unknown')}]\n"
         if event_type == LLM_TOOL_RESULT_EVENT:
-            return f"\nTool result: {json.dumps(data)}\n"
-        logger.error("Unknown event type: %s", event_type)
+            return "[Tool Result]\n"
+        if event_type == LLM_TURN_COMPLETE_EVENT:
+            return ""
         return ""
+
     return format_stream_data(
         {
             "event": event_type,

src/app/main.py

@@ -7,18 +7,18 @@
 from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
-from starlette.routing import Mount, Route, WebSocketRoute
 from llama_stack_client import APIConnectionError
+from starlette.routing import Mount, Route, WebSocketRoute
 
-from authorization.azure_token_manager import AzureEntraIDManager
 import metrics
 import version
+from a2a_storage import A2AStorageFactory
 from app import routers
 from app.database import create_tables, initialize_database
+from authorization.azure_token_manager import AzureEntraIDManager
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
 from log import get_logger
-from a2a_storage import A2AStorageFactory
 from models.responses import InternalServerErrorResponse
 from utils.common import register_mcp_servers_async
 from utils.llama_stack_version import check_llama_stack_version