from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import httpx

app = FastAPI()

class QueryRequest(BaseModel):
    query: str
    max_results: int = 10
    debug: bool = False

class QueryResponse(BaseModel):
    answer: str
    citations: list
    retrieval_trace: dict | None = None

@app.post("/query", response_model=QueryResponse)
async def query(request: QueryRequest):
    """
    Main query endpoint - forwards to orchestrator
    """
    async with httpx.AsyncClient() as client:
        response = await client.post(
            "http://orchestrator:8002/orchestrate",
            json=request.dict(),
            timeout=30.0
        )
        if response.status_code != 200:
            raise HTTPException(
                status_code=response.status_code,
                detail="Orchestrator error"
            )
        return response.json()

@app.get("/health")
async def health():
    """Health check endpoint"""
    return {"status": "healthy", "service": "api_gateway"}

from typing import List
from ..clients.vector_client import VectorClient
from ..clients.ontology_client import OntologyClient

class HopRetriever:
    """
    Executes multi-hop retrieval strategy
    """
    def __init__(self):
        self.vector_client = VectorClient()
        self.ontology_client = OntologyClient()

    async def retrieve(
        self,
        query: str,
        query_embedding: List[float],
        entities: List[str],
        temporal_constraint: dict | None = None
    ) -> dict:
        """
        Execute 3-hop retrieval
        """
        results = {"hops": []}

        # Hop 1: Dense vector search
        hop1_results = await self.vector_client.search(
            embedding=query_embedding,
            limit=20,
            metadata_filter=temporal_constraint
        )
        results["hops"].append({
            "hop": 1,
            "type": "vector_search",
            "count": len(hop1_results),
            "results": hop1_results
        })

        # Hop 2: Graph-guided retrieval
        # Get related entities from graph
        expanded_entities = await self.ontology_client.expand(entities)

        hop2_results = []
        for entity in expanded_entities:
            entity_docs = await self.vector_client.search_by_entity(
                entity=entity,
                embedding=query_embedding,
                limit=10
            )
            hop2_results.extend(entity_docs)

        results["hops"].append({
            "hop": 2,
            "type": "graph_guided",
            "expanded_entities": expanded_entities,
            "count": len(hop2_results),
            "results": hop2_results
        })

        # Hop 3: Exact match on metadata/keywords
        hop3_results = await self.vector_client.exact_match(
            entities=entities,
            temporal_constraint=temporal_constraint,
            limit=10
        )
        results["hops"].append({
            "hop": 3,
            "type": "exact_match",
            "count": len(hop3_results),
            "results": hop3_results
        })

        return results

from typing import List, Dict
import numpy as np

class FusionEngine:
    """
    Merges and reranks results from multiple hops
    """

    def fuse(
        self,
        hop_results: dict,
        query_info: dict,
        weights: dict = None
    ) -> List[dict]:
        """
        Combine results from all hops with multi-signal scoring
        """
        if weights is None:
            weights = {
                "semantic": 0.25,
                "graph": 0.30,
                "temporal": 0.25,
                "exact": 0.15,
                "hop": 0.05
            }

        # Collect all documents
        all_docs = []
        for hop in hop_results["hops"]:
            for doc in hop["results"]:
                doc["hop"] = hop["hop"]
                doc["hop_type"] = hop["type"]
                all_docs.append(doc)

        # Deduplicate by document ID
        unique_docs = self._deduplicate(all_docs)

        # Score each document
        scored_docs = []
        for doc in unique_docs:
            score = self._compute_fusion_score(doc, query_info, weights)
            doc["fusion_score"] = score
            scored_docs.append(doc)

        # Sort by score (descending)
        scored_docs.sort(key=lambda x: x["fusion_score"], reverse=True)

        return scored_docs

    def _compute_fusion_score(
        self,
        doc: dict,
        query_info: dict,
        weights: dict
    ) -> float:
        """
        Multi-signal scoring
        """
        scores = {
            "semantic": doc.get("similarity_score", 0.0),
            "graph": self._graph_score(doc, query_info),
            "temporal": self._temporal_score(doc, query_info),
            "exact": 1.0 if doc.get("exact_match") else 0.0,
            "hop": 1.0 if doc["hop"] == 1 else 0.9 if doc["hop"] == 2 else 0.8
        }

        final_score = sum(
            scores[key] * weights[key]
            for key in scores
        )
        return final_score

    def _graph_score(self, doc: dict, query_info: dict) -> float:
        """
        Score based on entity overlap with graph expansion
        """
        doc_entities = set(doc.get("entities", []))
        query_entities = set(query_info.get("expanded_entities", []))

        if not query_entities:
            return 0.5  # Neutral if no graph info

        overlap = len(doc_entities & query_entities)
        return min(overlap / len(query_entities), 1.0)

    def _temporal_score(self, doc: dict, query_info: dict) -> float:
        """
        Score based on temporal alignment
        """
        constraint = query_info.get("temporal_constraint")
        if not constraint:
            return 1.0  # No constraint = perfect score

        doc_day = doc.get("metadata", {}).get("day")
        if doc_day is None:
            return 0.5  # Unknown = neutral

        ref_day = constraint.get("reference")
        if constraint["type"] == "after":
            return 1.0 if doc_day > ref_day else 0.0
        elif constraint["type"] == "before":
            return 1.0 if doc_day < ref_day else 0.0
        elif constraint["type"] == "exact":
            return 1.0 if doc_day == ref_day else 0.0

        return 0.5

    def _deduplicate(self, docs: List[dict]) -> List[dict]:
        """Remove duplicate documents by ID"""
        seen = set()
        unique = []
        for doc in docs:
            if doc["id"] not in seen:
                seen.add(doc["id"])
                unique.append(doc)
        return unique

from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
from typing import List

app = FastAPI()

# Load model at startup
model = SentenceTransformer('all-MiniLM-L6-v2')

class EmbedRequest(BaseModel):
    text: str | List[str]

class EmbedResponse(BaseModel):
    embeddings: List[List[float]]
    model: str
    dimension: int

@app.post("/embed/query", response_model=EmbedResponse)
async def embed_query(request: EmbedRequest):
    """
    Embed a single query or multiple queries
    """
    texts = [request.text] if isinstance(request.text, str) else request.text

    embeddings = model.encode(
        texts,
        convert_to_numpy=True,
        show_progress_bar=False
    )

    return EmbedResponse(
        embeddings=embeddings.tolist(),
        model=model.get_config_dict()["name"],
        dimension=embeddings.shape[1]
    )

@app.post("/embed/document", response_model=EmbedResponse)
async def embed_document(request: EmbedRequest):
    """
    Embed document chunks (same implementation, different endpoint)
    """
    return await embed_query(request)

@app.post("/embed/batch")
async def embed_batch(texts: List[str]):
    """
    Batch embedding for pipeline processing
    """
    embeddings = model.encode(
        texts,
        batch_size=32,
        convert_to_numpy=True,
        show_progress_bar=True
    )
    return {
        "count": len(texts),
        "embeddings": embeddings.tolist()
    }

@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "service": "embeddings",
        "model": model.get_config_dict()["name"]
    }

import asyncpg
from typing import List

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

    async def connect(self):
        self.pool = await asyncpg.create_pool(
            host="vector_db",
            port=5432,
            database="ragdb",
            user="raguser",
            password="ragpass"
        )

    async def search(
        self,
        embedding: List[float],
        limit: int = 20,
        metadata_filter: dict | None = None
    ) -> List[dict]:
        """
        Vector similarity search with optional metadata filtering
        """
        query = """
            SELECT
                id,
                title,
                content,
                entities,
                metadata,
                1 - (embedding <=> $1) as similarity
            FROM documents
        """

        # Add metadata filter if provided
        if metadata_filter:
            if "day" in metadata_filter:
                query += " WHERE (metadata->>'day')::int > $3"

        query += " ORDER BY embedding <=> $1 LIMIT $2"

        async with self.pool.acquire() as conn:
            if metadata_filter and "day" in metadata_filter:
                rows = await conn.fetch(
                    query,
                    embedding,
                    limit,
                    metadata_filter["day"]
                )
            else:
                rows = await conn.fetch(query, embedding, limit)

        return [dict(row) for row in rows]

from fastapi import FastAPI
from pydantic import BaseModel
import networkx as nx
from typing import List

app = FastAPI()

# Load graph at startup
G = nx.DiGraph()

@app.on_event("startup")
async def load_graph():
    """Load graph from CSV triples"""
    # Format: subject,relation,object
    with open("src/graph/triples.csv") as f:
        for line in f:
            subj, rel, obj = line.strip().split(",")
            G.add_edge(subj, obj, relation=rel)

class ExpandRequest(BaseModel):
    entities: List[str]
    max_hops: int = 2
    relations: List[str] | None = None

@app.post("/ontology/expand")
async def expand(request: ExpandRequest):
    """
    Expand entities to their graph neighbors
    """
    expanded = set(request.entities)

    for entity in request.entities:
        if entity not in G:
            continue

        # Get neighbors within max_hops
        for _ in range(request.max_hops):
            neighbors = list(G.neighbors(entity))

            # Filter by relation type if specified
            if request.relations:
                neighbors = [
                    n for n in neighbors
                    if G[entity][n].get("relation") in request.relations
                ]

            expanded.update(neighbors)

    return {
        "original": request.entities,
        "expanded": list(expanded),
        "count": len(expanded)
    }

@app.post("/ontology/multihop")
async def multihop(start: str, end: str, max_depth: int = 3):
    """
    Find paths between two entities
    """
    try:
        paths = list(nx.all_simple_paths(
            G, start, end, cutoff=max_depth
        ))
        return {"paths": paths, "count": len(paths)}
    except nx.NodeNotFound:
        return {"paths": [], "count": 0}

@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "service": "ontology",
        "nodes": G.number_of_nodes(),
        "edges": G.number_of_edges()
    }

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import openai
import os

app = FastAPI()

openai.api_key = os.getenv("OPENAI_API_KEY")

class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 1000
    temperature: float = 0.7

class GenerateResponse(BaseModel):
    answer: str
    model: str
    tokens_used: int
    finish_reason: str

@app.post("/llm/generate", response_model=GenerateResponse)
async def generate(request: GenerateRequest):
    """
    Generate completion from LLM
    """
    try:
        response = await openai.ChatCompletion.acreate(
            model="gpt-4",
            messages=[
                {"role": "user", "content": request.prompt}
            ],
            max_tokens=request.max_tokens,
            temperature=request.temperature
        )

        return GenerateResponse(
            answer=response.choices[0].message.content,
            model=response.model,
            tokens_used=response.usage.total_tokens,
            finish_reason=response.choices[0].finish_reason
        )

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"LLM generation failed: {str(e)}"
        )

@app.get("/health")
async def health():
    return {"status": "healthy", "service": "llm_proxy"}

from fastapi import FastAPI
import asyncio

app = FastAPI()

@app.post("/pipeline/run")
async def run_pipeline(source: str, domain: str = "default"):
    """
    Run full ingestion pipeline
    """
    steps = [
        ingest_documents(source),
        clean_and_chunk(),
        extract_entities(),
        build_graph_triples(),
        generate_embeddings(),
        load_to_databases()
    ]

    results = {"steps": []}
    for step in steps:
        result = await step
        results["steps"].append(result)

    return results

async def ingest_documents(source: str):
    """Load raw documents"""
    # Implementation here
    return {"step": "ingest", "status": "success"}

async def extract_entities():
    """Extract entities using NER"""
    # Use spaCy or similar
    return {"step": "entities", "status": "success"}

services:
  orchestrator:
    # Can reach other services by name
    # http://embeddings:8003
    # http://ontology_service:8004
    depends_on:
      - embeddings
      - ontology_service
      - vector_db