import pytest
from orchestrator.services.query_parser import QueryParser

def test_entity_extraction():
    parser = QueryParser()
    result = parser.parse("Who helped Arjuna after Day 10 of the war?")

    assert "Arjuna" in result.entities
    assert result.query_type == "temporal_causal"
    assert result.temporal_constraint["type"] == "after"
    assert result.temporal_constraint["reference"] == 10

def test_intent_classification():
    parser = QueryParser()

    # Test who question
    result1 = parser.parse("Who was Arjuna's ally?")
    assert result1.intent == "relational"
    assert result1.query_type == "who_question"

    # Test what question
    result2 = parser.parse("What happened on Day 10?")
    assert result2.intent == "factual"
    assert result2.query_type == "what_question"

    # Test why question
    result3 = parser.parse("Why did Krishna advise Arjuna?")
    assert result3.intent == "causal"
    assert result3.query_type == "why_question"

import pytest
import asyncio
from orchestrator.clients.vector_client import VectorClient

@pytest.mark.integration
@pytest.mark.asyncio
async def test_vector_search_integration():
    """
    Test actual connection to vector DB
    Requires: docker compose up vector_db
    """
    client = VectorClient()
    await client.connect()

    # Create test embedding
    test_embedding = [0.1] * 384  # Match model dimension

    # Search
    results = await client.search(
        embedding=test_embedding,
        limit=10
    )

    assert len(results) > 0
    assert "id" in results[0]
    assert "content" in results[0]
    assert "similarity" in results[0]

@pytest.mark.integration
@pytest.mark.asyncio
async def test_ontology_expansion_integration():
    """Test ontology service integration"""
    from orchestrator.clients.ontology_client import OntologyClient

    client = OntologyClient("http://ontology_service:8004")

    result = await client.expand(
        entities=["Arjuna"],
        max_hops=2
    )

    assert "expanded" in result
    assert "Krishna" in result["expanded"]  # Known ally
    assert "Bhima" in result["expanded"]    # Known brother

import pytest
import httpx

@pytest.mark.e2e
@pytest.mark.asyncio
async def test_complete_query_flow():
    """
    Test complete query from API gateway to answer
    Requires: All services running (docker compose up)
    """
    async with httpx.AsyncClient() as client:
        response = await client.post(
            "http://localhost:8001/query",
            json={
                "query": "Who helped Arjuna after Day 10?",
                "max_results": 5,
                "debug": True
            },
            timeout=30.0
        )

    assert response.status_code == 200
    data = response.json()

    # Validate response structure
    assert "answer" in data
    assert "citations" in data
    assert len(data["citations"]) > 0

    # Validate answer quality
    answer_lower = data["answer"].lower()
    assert any(name in answer_lower for name in ["krishna", "bhima"])

    # Validate debug trace
    if "retrieval_trace" in data:
        trace = data["retrieval_trace"]
        assert "hops" in trace
        assert len(trace["hops"]) >= 2  # At least 2 hops

import pytest
from evaluation.metrics import (
    answer_relevance,
    citation_precision,
    citation_recall
)

@pytest.mark.evaluation
def test_answer_quality():
    """
    Evaluate system on benchmark questions
    """
    test_cases = [
        {
            "query": "Who helped Arjuna after Day 10?",
            "expected_entities": ["Krishna", "Bhima"],
            "expected_citations": ["doc_156", "doc_203"]
        },
        # More test cases...
    ]

    results = []
    for case in test_cases:
        response = query_system(case["query"])

        # Check if expected entities appear in answer
        relevance = answer_relevance(
            response["answer"],
            case["expected_entities"]
        )

        # Check citation quality
        precision = citation_precision(
            response["citations"],
            case["expected_citations"]
        )

        results.append({
            "query": case["query"],
            "relevance": relevance,
            "precision": precision
        })

    # Assert minimum quality
    avg_relevance = sum(r["relevance"] for r in results) / len(results)
    assert avg_relevance > 0.8, f"Low answer relevance: {avg_relevance}"

# VPC and Networking
resource "aws_vpc" "ragforge" {
  cidr_block = "10.0.0.0/16"
  enable_dns_hostnames = true
  enable_dns_support = true

  tags = {
    Name = "ragforge-vpc"
  }
}

# RDS PostgreSQL with pgvector
resource "aws_db_instance" "vector_db" {
  identifier = "ragforge-vector-db"
  engine = "postgres"
  engine_version = "16.1"
  instance_class = "db.t3.large"

  allocated_storage = 100
  storage_type = "gp3"

  db_name = "ragdb"
  username = "raguser"
  password = data.aws_secretsmanager_secret_version.db_password.secret_string

  vpc_security_group_ids = [aws_security_group.db.id]
  db_subnet_group_name = aws_db_subnet_group.ragforge.name

  # Enable pgvector extension via custom parameter group
  parameter_group_name = aws_db_parameter_group.pgvector.name

  backup_retention_period = 7
  skip_final_snapshot = false
  final_snapshot_identifier = "ragforge-final-snapshot"

  tags = {
    Name = "ragforge-vector-db"
  }
}

# ECS Cluster
resource "aws_ecs_cluster" "ragforge" {
  name = "ragforge-cluster"

  setting {
    name = "containerInsights"
    value = "enabled"
  }
}

# ECS Task Definition - Orchestrator
resource "aws_ecs_task_definition" "orchestrator" {
  family = "ragforge-orchestrator"
  requires_compatibilities = ["FARGATE"]
  network_mode = "awsvpc"
  cpu = "1024"  # 1 vCPU
  memory = "2048"  # 2 GB

  container_definitions = jsonencode([
    {
      name = "orchestrator"
      image = "${var.ecr_registry}/ragforge-orchestrator:latest"
      portMappings = [
        {
          containerPort = 8002
          protocol = "tcp"
        }
      ]
      environment = [
        {
          name = "POSTGRES_HOST"
          value = aws_db_instance.vector_db.endpoint
        },
        {
          name = "EMBEDDING_SERVICE_URL"
          value = "http://${aws_service_discovery_service.embeddings.name}.${aws_service_discovery_private_dns_namespace.ragforge.name}:8003"
        }
      ]
      secrets = [
        {
          name = "POSTGRES_PASSWORD"
          valueFrom = aws_secretsmanager_secret.db_password.arn
        }
      ]
      logConfiguration = {
        logDriver = "awslogs"
        options = {
          "awslogs-group" = "/ecs/ragforge-orchestrator"
          "awslogs-region" = var.aws_region
          "awslogs-stream-prefix" = "ecs"
        }
      }
    }
  ])
}

# ECS Service - Orchestrator
resource "aws_ecs_service" "orchestrator" {
  name = "orchestrator"
  cluster = aws_ecs_cluster.ragforge.id
  task_definition = aws_ecs_task_definition.orchestrator.arn
  desired_count = 2
  launch_type = "FARGATE"

  network_configuration {
    subnets = aws_subnet.private[*].id
    security_groups = [aws_security_group.ecs_tasks.id]
  }

  load_balancer {
    target_group_arn = aws_lb_target_group.orchestrator.arn
    container_name = "orchestrator"
    container_port = 8002
  }

  service_registries {
    registry_arn = aws_service_discovery_service.orchestrator.arn
  }

  depends_on = [aws_lb_listener.api]
}

# Similar definitions for other services...

import uuid
from contextvars import ContextVar

# Context variable for trace ID
trace_id_var: ContextVar[str] = ContextVar('trace_id', default='')

def get_trace_id() -> str:
    return trace_id_var.get()

def set_trace_id(trace_id: str):
    trace_id_var.set(trace_id)

# In FastAPI middleware
from starlette.middleware.base import BaseHTTPMiddleware

class TraceMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        # Get or generate trace ID
        trace_id = request.headers.get('X-Trace-ID', str(uuid.uuid4()))
        set_trace_id(trace_id)

        # Add to response headers
        response = await call_next(request)
        response.headers['X-Trace-ID'] = trace_id
        return response

# Add to app
app.add_middleware(TraceMiddleware)

import logging
import json
from .observability import get_trace_id

class StructuredLogger:
    def __init__(self, name: str):
        self.logger = logging.getLogger(name)
        self.logger.setLevel(logging.INFO)

    def info(self, message: str, **kwargs):
        log_data = {
            "level": "INFO",
            "message": message,
            "trace_id": get_trace_id(),
            "service": "orchestrator",
            **kwargs
        }
        self.logger.info(json.dumps(log_data))

# Usage
logger = StructuredLogger(__name__)

logger.info(
    "Starting multi-hop retrieval",
    query=query,
    entities=entities,
    hops=3
)

from prometheus_client import Counter, Histogram
import time

# Define metrics
query_counter = Counter(
    'ragforge_queries_total',
    'Total number of queries',
    ['query_type', 'status']
)

query_latency = Histogram(
    'ragforge_query_latency_seconds',
    'Query latency in seconds',
    ['service', 'operation']
)

# Usage
@app.post("/orchestrate")
async def orchestrate(request: QueryRequest):
    start_time = time.time()

    try:
        result = await do_orchestration(request)

        # Record success
        query_counter.labels(
            query_type=result.query_type,
            status="success"
        ).inc()

        return result

    except Exception as e:
        # Record failure
        query_counter.labels(
            query_type="unknown",
            status="error"
        ).inc()
        raise

    finally:
        # Record latency
        latency = time.time() - start_time
        query_latency.labels(
            service="orchestrator",
            operation="full_query"
        ).observe(latency)

@app.post("/debug/orchestrate")
async def debug_orchestrate(request: QueryRequest):
    """
    Returns verbose debugging information about retrieval
    """
    trace = {
        "trace_id": get_trace_id(),
        "steps": []
    }

    # Step 1: Parse query
    start = time.time()
    parsed = await query_parser.parse(request.query)
    trace["steps"].append({
        "step": "parse_query",
        "latency_ms": (time.time() - start) * 1000,
        "result": parsed.dict()
    })

    # Step 2: Semantic expansion
    start = time.time()
    expanded = await semantic_expander.expand(parsed)
    trace["steps"].append({
        "step": "semantic_expansion",
        "latency_ms": (time.time() - start) * 1000,
        "expansions": expanded
    })

    # ... more steps with timing ...

    return {
        "answer": final_answer,
        "debug_trace": trace
    }

from functools import lru_cache
import redis
import pickle

# Redis for distributed caching
redis_client = redis.Redis(host='redis', port=6379)

def cache_embedding(text: str, embedding: list):
    """Cache query embeddings"""
    key = f"embedding:{hash(text)}"
    redis_client.setex(
        key,
        3600,  # 1 hour TTL
        pickle.dumps(embedding)
    )

def get_cached_embedding(text: str):
    """Retrieve cached embedding"""
    key = f"embedding:{hash(text)}"
    cached = redis_client.get(key)
    if cached:
        return pickle.loads(cached)
    return None

# In-memory cache for ontology expansions
@lru_cache(maxsize=1000)
def expand_entity(entity: str, max_hops: int):
    """Cache common entity expansions"""
    return ontology_client.expand([entity], max_hops)

import asyncpg

class VectorClient:
    def __init__(self):
        self.pool = None

    async def initialize(self):
        """Create connection pool at startup"""
        self.pool = await asyncpg.create_pool(
            host=os.getenv("POSTGRES_HOST"),
            port=5432,
            database="ragdb",
            user="raguser",
            password=os.getenv("POSTGRES_PASSWORD"),
            min_size=5,
            max_size=20
        )

    async def search(self, embedding, limit=20):
        async with self.pool.acquire() as conn:
            # Use pooled connection
            results = await conn.fetch(query, embedding, limit)
        return results

import asyncio

async def retrieve_multi_hop(query_info):
    """Execute hops in parallel where possible"""

    # Get embedding and ontology expansion in parallel
    embedding_task = embeddings_client.embed(query_info.query)
    ontology_task = ontology_client.expand(query_info.entities)

    embedding, expanded_entities = await asyncio.gather(
        embedding_task,
        ontology_task
    )

    # Now execute multiple Hop 2 queries in parallel
    hop2_tasks = [
        vector_client.search_by_entity(entity, embedding)
        for entity in expanded_entities
    ]

    hop2_results = await asyncio.gather(*hop2_tasks)

    return combine_results(hop2_results)

#   - pubmed_abstracts.json
#   - clinical_trials.json
#   - drug_interactions.csv

ENTITY_TYPES = [
    "Drug",
    "Disease",
    "Symptom",
    "Treatment",
    "Gene",
    "Protein"
]

RELATIONS = [
    "TREATS",         # Drug -[TREATS]-> Disease
    "CAUSES",         # Gene -[CAUSES]-> Disease
    "HAS_SYMPTOM",    # Disease -[HAS_SYMPTOM]-> Symptom
    "INTERACTS_WITH", # Drug -[INTERACTS_WITH]-> Drug
    "INHIBITS",       # Drug -[INHIBITS]-> Protein
]

import spacy

# Load medical NER model
nlp = spacy.load("en_ner_bc5cdr_md")  # Bio-medical NER

def extract_medical_entities(text: str):
    doc = nlp(text)

    entities = []
    for ent in doc.ents:
        entities.append({
            "text": ent.text,
            "type": ent.label_,  # DISEASE, CHEMICAL
            "start": ent.start_char,
            "end": ent.end_char
        })

    return entities