from sentence_transformers import CrossEncoder

class CrossEncoderReranker:
    def __init__(self):
        self.model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

    def rerank(self, query: str, results: list[SearchResult]) -> list[SearchResult]:
        pairs = [(query, r.content) for r in results]
        scores = self.model.predict(pairs)

        scored_results = list(zip(results, scores))
        scored_results.sort(key=lambda x: x[1], reverse=True)

        return [r for r, _ in scored_results[:TOP_K]]

import litellm

class LiteLLMRouter:
    """Unified interface for multiple LLM providers."""

    def __init__(self, model_list: list[dict]):
        self.router = litellm.Router(model_list=model_list)

    async def generate(self, messages: list[dict], **kwargs) -> str:
        response = await self.router.acompletion(
            model="gpt-4o",  # Routes based on availability
            messages=messages,
            **kwargs
        )
        return response.choices[0].message.content

from fastapi.responses import StreamingResponse

@app.post("/api/explain")
async def explain_post(request: ExplainRequest):
    async def event_stream():
        # Stage 1: Extract post
        post = await extractor.extract(request.url)
        yield f"event: post\ndata: {post.json()}\n\n"

        # Stage 2: Search for context
        sources = await searcher.search(post.text)
        yield f"event: sources\ndata: {json.dumps(sources)}\n\n"

        # Stage 3: Rerank results
        ranked = reranker.rerank(post.text, sources)

        # Stage 4: Generate explanation bullets
        async for bullet in generator.stream_bullets(post, ranked):
            yield f"event: bullet\ndata: {bullet.json()}\n\n"

        yield f"event: done\ndata: {{}}\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")