# Story 5.1: Query Endpoint

Status: done

## Story

As an **API consumer**,
I want **to submit natural language questions**,
So that **I can get answers from my uploaded documents**.

## Acceptance Criteria

1. **AC5.1.1:** POST /query endpoint accepts natural language queries
   - Request body includes `query` (string, required)
   - Request body includes `document_ids` (list[UUID], optional) - filters search to specific documents
   - Request body includes `include_context` (bool, optional, default=true) - includes source references in response
   - Validates query is non-empty and ≤ 1000 characters
   - Returns 400 for invalid requests
   - Returns 401 if API key authentication fails

2. **AC5.1.2:** Response includes answer with metadata
   - Returns JSON with `query`, `answer`, `confidence`, `sources`, `processing_time_ms`
   - `answer` is a natural language response generated from retrieved content
   - `confidence` is a float between 0.0 and 1.0 indicating answer quality
   - `sources` is a list of source references with file, sheet, location, context
   - `processing_time_ms` is the total time taken to process the query
   - Returns 200 for successful queries

3. **AC5.1.3:** Source references include document context
   - Each source includes: `file`, `sheet`, `location` (cell range), `context` (snippet)
   - Context is the actual text content from the extracted record
   - Location is the Excel cell range (e.g., "A42:D42")
   - Sheet name corresponds to the original Excel sheet
   - Sources are ordered by relevance (highest similarity first)

4. **AC5.1.4:** Optional document_ids filter
   - If `document_ids` is provided, search only those documents
   - If `document_ids` is null/omitted, search all processed documents
   - Return 404 if any document_id doesn't exist or isn't processed
   - Filter applies to both retrieval and source references

5. **AC5.1.5:** Performance constraint
   - Total response time ≤ 5 seconds for typical queries
   - Log processing time for monitoring
   - Return partial results if LLM generation exceeds time limit
   - Use timeout protection for external API calls (Azure OpenAI)

## Tasks / Subtasks

- [ ] **Task 1: Create query API route** (AC: 5.1.1)
  - [ ] Create `src/api/routes/query.py` with POST /query endpoint
  - [ ] Define QueryRequest Pydantic model (query, document_ids, include_context)
  - [ ] Define QueryResponse Pydantic model (query, answer, confidence, sources, processing_time_ms)
  - [ ] Define SourceReference Pydantic model (file, sheet, location, context)
  - [ ] Add input validation for query length and format
  - [ ] Integrate API key authentication middleware
  - [ ] Add structured logging for all query requests

- [ ] **Task 2: Implement semantic search retrieval** (AC: 5.1.2, 5.1.3, 5.1.4)
  - [ ] Create `src/retrieval/semantic_retriever.py` service
  - [ ] Accept query string and optional document_ids filter
  - [ ] Generate query embedding using EmbeddingService
  - [ ] Search Milvus with COSINE similarity
  - [ ] Apply document_ids filter using Milvus expression
  - [ ] Retrieve top-k results (default k=5) with metadata
  - [ ] Format results as SourceReference objects
  - [ ] Return empty list if no results found

- [ ] **Task 3: Implement answer generation** (AC: 5.1.2, 5.1.3)
  - [ ] Create `src/retrieval/answer_generator.py` service
  - [ ] Accept query and list of retrieved sources
  - [ ] Format sources as context for LLM prompt
  - [ ] Call Azure OpenAI (GPT-4) with query + context
  - [ ] Parse LLM response to extract answer
  - [ ] Calculate confidence score based on retrieval similarity
  - [ ] Handle cases with no retrieved sources (return low confidence)
  - [ ] Add retry logic with exponential backoff for API failures

- [ ] **Task 4: Integrate retrieval and generation in endpoint** (AC: 5.1.1, 5.1.2, 5.1.5)
  - [ ] Wire SemanticRetriever and AnswerGenerator in query endpoint
  - [ ] Track processing time from request start to response
  - [ ] Validate document_ids exist in database before querying
  - [ ] Return 404 if document_ids are invalid or not processed
  - [ ] Add timeout protection (5 seconds max)
  - [ ] Return partial results if timeout is approaching
  - [ ] Log all errors with query_id for debugging

- [ ] **Task 5: Add query endpoint to FastAPI app** (AC: 5.1.1)
  - [ ] Register query router in `src/api/main.py`
  - [ ] Add /query route to API prefix structure
  - [ ] Update OpenAPI documentation with query examples
  - [ ] Add CORS configuration if needed
  - [ ] Ensure authentication middleware is applied

- [ ] **Task 6: Unit tests for retrieval and generation** (AC: All)
  - [ ] Test SemanticRetriever with mock Milvus
  - [ ] Test AnswerGenerator with mock Azure OpenAI
  - [ ] Test document_ids filtering logic
  - [ ] Test error handling for missing documents
  - [ ] Test timeout handling
  - [ ] Mock embeddings and search results

- [ ] **Task 7: Integration tests for /query endpoint** (AC: All)
  - [ ] Test POST /query with valid query (returns answer)
  - [ ] Test POST /query with document_ids filter
  - [ ] Test POST /query without document_ids (searches all)
  - [ ] Test invalid query (empty, too long)
  - [ ] Test invalid document_ids (non-existent UUIDs)
  - [ ] Test authentication (valid API key, missing API key)
  - [ ] Test response time constraint
  - [ ] Use real PostgreSQL, Milvus, and Azure OpenAI

- [ ] **Task 8: Create examples and update documentation** (AC: All)
  - [ ] Create `examples/test_query_endpoint.py` with usage examples
  - [ ] Demonstrate query with and without document_ids filter
  - [ ] Show source reference parsing
  - [ ] Update README with /query endpoint documentation
  - [ ] Add query examples to API documentation

## Dev Notes

### Architecture Patterns and Constraints

**From architecture.md:**
- **API Layer:** Create `src/api/routes/query.py` for query endpoint
- **Service Layer:** Create `src/retrieval/semantic_retriever.py` and `src/retrieval/answer_generator.py`
- **Authentication:** Use API key middleware from Story 2.1
- **Error Handling:** Use `QueryError`, `DocumentNotFoundError` from `src.core.exceptions`
- **Logging:** Structured logging with structlog (log query_id, processing_time, document_ids)
- **Models:** Pydantic models for request/response validation

**From PRD:**
- **FR8:** Users can query documents with natural language questions
- **FR9:** System returns answers with source citations
- **NFR2:** Response time ≤ 5 seconds for queries
- **NFR7:** System uses Azure OpenAI for embeddings and answer generation

### Learnings from Previous Story

**From Story 4.5: Incremental Indexing (Status: done)**

**Services Created:**
- `VectorStore` at `src/knowledge/vector_store.py` - Use for semantic search
- `EmbeddingService` at `src/knowledge/embeddings.py` - Use for query embeddings
- `DocumentDeletionService` - Pattern for orchestrating multiple stores

**Key Patterns:**
- **Async operations:** Use async/await throughout for non-blocking I/O
- **Retry logic:** Exponential backoff (1s, 2s, 4s) for Azure OpenAI API calls
- **Structured logging:** Log all operations with contextual fields (query_id, document_ids, processing_time)
- **Error handling:** Catch and log all exceptions, return appropriate HTTP status codes
- **Milvus search:** Use COSINE similarity with HNSW index for vector search

**Technical Details:**
- **Milvus collection:** `dsol_records` with 3072-dim vectors
- **Distance metric:** COSINE similarity (higher = more similar)
- **Embedding model:** Azure OpenAI text-embedding-3-large
- **Search parameters:** `{"metric_type": "COSINE", "params": {"ef": 64}}`
- **Top-k results:** Default to 5 for query retrieval

**Integration Points:**
- Use `EmbeddingService.embed_batch([query])` to generate query embedding
- Use Milvus `collection.search()` with document_id filter expression
- Retrieve `text_content`, `filename`, `sheet_name`, `row_number`, `col_range` fields
- Format results as source references with proper metadata

[Source: docs/sprint-artifacts/4-5-incremental-indexing.md]

### Milvus Semantic Search Pattern

**Query Embedding:**
```python
from src.knowledge.embeddings import EmbeddingService

embedding_service = EmbeddingService()
query_embedding = await embedding_service.embed_batch([query_text])
```

**Milvus Search:**
```python
from pymilvus import Collection, connections
from src.core.config import settings

connections.connect(host=settings.milvus_host, port=settings.milvus_port)
collection = Collection(settings.milvus_collection)
collection.load()

# Build filter expression
expr = None
if document_ids:
    uuid_strs = [f'"{str(doc_id)}"' for doc_id in document_ids]
    expr = f'document_id in [{", ".join(uuid_strs)}]'

# Search
results = collection.search(
    data=query_embedding,
    anns_field="vector",
    param={"metric_type": "COSINE", "params": {"ef": 64}},
    limit=5,
    expr=expr,
    output_fields=["text_content", "filename", "sheet_name", "row_number", "col_range"]
)

# Extract results
sources = []
for hit in results[0]:
    sources.append({
        "file": hit.entity.get("filename"),
        "sheet": hit.entity.get("sheet_name"),
        "location": hit.entity.get("col_range"),
        "context": hit.entity.get("text_content"),
        "similarity": hit.distance
    })
```

### Azure OpenAI Answer Generation Pattern

**LLM Call with Context:**
```python
from openai import AsyncAzureOpenAI
from src.core.config import settings

client = AsyncAzureOpenAI(
    api_key=settings.azure_openai_api_key,
    api_version=settings.azure_openai_api_version,
    azure_endpoint=settings.azure_openai_endpoint
)

# Format context from sources
context = "\n\n".join([
    f"Source {i+1} ({src['file']}, {src['sheet']}, {src['location']}):\n{src['context']}"
    for i, src in enumerate(sources)
])

# Build prompt
system_prompt = """You are a helpful assistant that answers questions based on provided document context.
Always cite your sources and indicate confidence in your answer."""

user_prompt = f"""Context from documents:
{context}

Question: {query}

Please provide a clear, concise answer based only on the context provided."""

# Call LLM
response = await client.chat.completions.create(
    model=settings.azure_openai_chat_deployment,
    messages=[
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt}
    ],
    temperature=0.3,
    max_tokens=500
)

answer = response.choices[0].message.content
```

### Confidence Score Calculation

**Simple heuristic based on retrieval similarity:**
```python
def calculate_confidence(sources: list[dict]) -> float:
    """Calculate confidence score based on retrieval results.

    - No sources: 0.0
    - Low similarity (< 0.3): 0.3-0.5
    - Medium similarity (0.3-0.6): 0.5-0.7
    - High similarity (> 0.6): 0.7-1.0
    """
    if not sources:
        return 0.0

    # Use top result similarity
    top_similarity = sources[0].get("similarity", 0.0)

    if top_similarity < 0.3:
        return 0.3 + (top_similarity / 0.3) * 0.2  # 0.3-0.5
    elif top_similarity < 0.6:
        return 0.5 + ((top_similarity - 0.3) / 0.3) * 0.2  # 0.5-0.7
    else:
        return 0.7 + ((top_similarity - 0.6) / 0.4) * 0.3  # 0.7-1.0
```

### API Request/Response Models

**Request Model:**
```python
from pydantic import BaseModel, Field, validator
from uuid import UUID

class QueryRequest(BaseModel):
    query: str = Field(..., min_length=1, max_length=1000, description="Natural language query")
    document_ids: list[UUID] | None = Field(None, description="Optional document filter")
    include_context: bool = Field(True, description="Include source references in response")

    @validator("query")
    def query_not_empty(cls, v):
        if not v.strip():
            raise ValueError("Query cannot be empty")
        return v.strip()
```

**Response Model:**
```python
from pydantic import BaseModel

class SourceReference(BaseModel):
    file: str
    sheet: str
    location: str
    context: str

class QueryResponse(BaseModel):
    query: str
    answer: str
    confidence: float = Field(..., ge=0.0, le=1.0)
    sources: list[SourceReference]
    processing_time_ms: int
```

### Error Handling

**Document Validation:**
```python
from src.db.models import Document
from sqlalchemy import select

async def validate_document_ids(document_ids: list[UUID], session) -> None:
    """Validate that all document_ids exist and are processed.

    Raises:
        DocumentNotFoundError: If any document doesn't exist or isn't completed
    """
    for doc_id in document_ids:
        result = await session.execute(
            select(Document).where(
                Document.id == doc_id,
                Document.status == "completed"
            )
        )
        doc = result.scalar_one_or_none()
        if not doc:
            raise DocumentNotFoundError(
                f"Document {doc_id} not found or not processed"
            )
```

**Timeout Protection:**
```python
import asyncio
from fastapi import HTTPException

async def query_with_timeout(query: QueryRequest) -> QueryResponse:
    """Execute query with 5-second timeout."""
    try:
        return await asyncio.wait_for(
            process_query(query),
            timeout=5.0
        )
    except asyncio.TimeoutError:
        logger.warning("query_timeout", query=query.query)
        raise HTTPException(
            status_code=504,
            detail="Query processing timed out (>5 seconds)"
        )
```

### Integration with Document Lifecycle

**Document Status Check:**
- Only query documents with `status="completed"`
- Documents with `status="processing"` or `status="failed"` are excluded
- Use database query to filter by status before Milvus search

**Example Query Flow:**
```
1. Receive POST /query request
2. Validate API key (authentication middleware)
3. Validate request body (Pydantic)
4. Validate document_ids exist and are completed (database query)
5. Generate query embedding (Azure OpenAI)
6. Search Milvus with document_ids filter
7. Format sources from search results
8. Generate answer from LLM (Azure OpenAI + context)
9. Calculate confidence score
10. Track processing time
11. Return response
```

### Prerequisites

- **Story 4.4 COMPLETE:** VectorStore with Milvus indexing
- **Story 4.5 COMPLETE:** Incremental indexing and dual-store consistency
- **Story 2.1 COMPLETE:** API key authentication middleware
- **Database:** PostgreSQL with processed documents
- **Milvus:** Collection with indexed vectors
- **Azure OpenAI:** Configured endpoints for embeddings and chat

### Success Criteria

- ✅ POST /query endpoint accepts natural language queries
- ✅ Endpoint returns answers with confidence scores
- ✅ Source references include file, sheet, location, context
- ✅ Optional document_ids filter works correctly
- ✅ Response time ≤ 5 seconds for typical queries
- ✅ Authentication required (API key)
- ✅ Unit tests pass (10+ tests)
- ✅ Integration tests pass with real services
- ✅ Examples demonstrate query usage
- ✅ API documentation updated

### Out of Scope

- **Query routing:** Classification of query types (Story 5.2)
- **Hybrid search:** Combining exact and semantic search (Story 5.2)
- **Batch queries:** Multiple queries in single request (Story 5.5)
- **Query logging:** Persisting queries for analytics (Story 5.6)
- **Answer quality scoring:** Advanced confidence metrics beyond similarity
- **Caching:** Query result caching for performance

### Next Steps

After Story 5.1:
- **Story 5.2:** Query router (classify exact vs semantic queries)
- **Story 5.3:** Advanced answer generation with source citations
- **Story 5.4:** Confidence scoring improvements
- **Story 5.5:** Batch query support
- **Story 5.6:** Query logging and analytics

## Technical Design

### Data Flow Diagram

**Query Processing Flow:**
```
POST /query
    ↓
[Authentication] → Validate API Key
    ↓
[Validation] → Validate Request Body
    ↓
[Document Check] → Verify document_ids exist and are completed
    ↓
[Embedding] → Generate query embedding (Azure OpenAI)
    ↓
[Retrieval] → Search Milvus with document_ids filter → Top-k results
    ↓
[Formatting] → Format results as SourceReference objects
    ↓
[Generation] → Call LLM with query + context → Answer
    ↓
[Confidence] → Calculate confidence from similarity scores
    ↓
[Response] → Return QueryResponse with sources
```

### Module Structure

```
src/
├── api/
│   └── routes/
│       └── query.py              # NEW - POST /query endpoint
├── retrieval/                    # NEW module
│   ├── __init__.py
│   ├── semantic_retriever.py    # NEW - Milvus search service
│   └── answer_generator.py      # NEW - LLM answer generation
└── knowledge/
    ├── embeddings.py             # EXISTING - Use for query embeddings
    └── vector_store.py           # EXISTING - Use for Milvus connection
```

### API Response Examples

**Successful Query:**
```json
{
  "query": "What does item code 2024-4_0019 appear in?",
  "answer": "Item code 2024-4_0019 appears in the Screen column with status 'Applicable'. It is found in the Distribution Item Specification document in the Item Mapping sheet.",
  "confidence": 0.92,
  "sources": [
    {
      "file": "distribution_spec.xlsx",
      "sheet": "Item Mapping",
      "location": "A42:D42",
      "context": "Item Code: 2024-4_0019. Item Name: Distribution Module A. Remarks: Hardware integration component. Status: Active"
    }
  ],
  "processing_time_ms": 1250
}
```

**Query with document_ids Filter:**
```json
{
  "query": "What hardware components are listed?",
  "answer": "The document lists Server Hardware Package as a hardware component with category 'Hardware' priced at 5000.",
  "confidence": 0.85,
  "sources": [
    {
      "file": "inventory_2024.xlsx",
      "sheet": "Hardware Items",
      "location": "A5:D5",
      "context": "Item Code: HW-2024-001. Item Name: Server Hardware Package. Category: Hardware. Price: 5000"
    }
  ],
  "processing_time_ms": 980
}
```

**No Results Found:**
```json
{
  "query": "What is the weather today?",
  "answer": "I couldn't find any information related to your question in the uploaded documents.",
  "confidence": 0.1,
  "sources": [],
  "processing_time_ms": 450
}
```

**Error Response (Invalid Document ID):**
```json
{
  "detail": "Document 550e8400-e29b-41d4-a716-446655440000 not found or not processed"
}
```

### Example Usage

```python
import httpx
import asyncio

API_BASE_URL = "http://localhost:8000"
API_KEY = "your-api-key-here"

async def query_documents():
    """Example: Query documents with natural language."""
    async with httpx.AsyncClient() as client:
        # Query all documents
        response = await client.post(
            f"{API_BASE_URL}/query",
            json={
                "query": "What does item code 2024-4_0019 appear in?",
                "include_context": True
            },
            headers={"X-API-Key": API_KEY}
        )

        result = response.json()
        print(f"Answer: {result['answer']}")
        print(f"Confidence: {result['confidence']}")
        print(f"Sources: {len(result['sources'])}")
        print(f"Processing time: {result['processing_time_ms']}ms")

        # Query specific documents only
        response = await client.post(
            f"{API_BASE_URL}/query",
            json={
                "query": "What hardware components are available?",
                "document_ids": ["550e8400-e29b-41d4-a716-446655440000"],
                "include_context": True
            },
            headers={"X-API-Key": API_KEY}
        )

        result = response.json()
        print(f"\nFiltered query answer: {result['answer']}")

if __name__ == "__main__":
    asyncio.run(query_documents())
```

---

## Definition of Done

- [ ] All acceptance criteria (AC5.1.1 - AC5.1.5) implemented and tested
- [ ] POST /query endpoint created in `src/api/routes/query.py`
- [ ] SemanticRetriever service created in `src/retrieval/semantic_retriever.py`
- [ ] AnswerGenerator service created in `src/retrieval/answer_generator.py`
- [ ] QueryRequest and QueryResponse Pydantic models defined
- [ ] API key authentication applied to endpoint
- [ ] Document validation (only completed documents queryable)
- [ ] Semantic search with Milvus working (COSINE similarity)
- [ ] Query embedding generation using Azure OpenAI
- [ ] Answer generation using Azure OpenAI (GPT-4)
- [ ] Confidence score calculation implemented
- [ ] Source references formatted correctly
- [ ] document_ids filter working (optional)
- [ ] Processing time tracking implemented
- [ ] Response time ≤ 5 seconds enforced with timeout
- [ ] Unit tests pass (10+ tests)
- [ ] Integration tests pass with real PostgreSQL, Milvus, Azure OpenAI
- [ ] Error handling for invalid queries, missing documents, timeouts
- [ ] Examples created in `examples/test_query_endpoint.py`
- [ ] Code follows project style (Black, Ruff, mypy strict)
- [ ] Structured logging implemented (query_id, processing_time, document_ids)
- [ ] Documentation strings (docstrings) for all public methods
- [ ] API documentation updated (OpenAPI/Swagger)
- [ ] Sprint status updated to "done"
- [ ] Git commit with proper commit message

## Dev Agent Record

### Context Reference

<!-- Path(s) to story context XML will be added here by context workflow -->

### Agent Model Used

{{agent_model_name_version}}

### Debug Log References

### Completion Notes List

### File List