from unstructured.partition.pdf import partition_pdf

elements = partition_pdf(
    filename="technical-manual.pdf",
    strategy="hi_res",
    infer_table_structure=True,
    languages=["eng"],
    extract_images_in_pdf=True,
    extract_image_block_types=["Image", "Table"],
)

# Elements are typed: NarrativeText, Title, Table, ListItem, etc.
for element in elements:
    print(f"Type: {type(element).__name__}, Text: {element.text[:100]}")

# DO NOT use this in production
chunks = [text[i:i+512] for i in range(0, len(text), 512)]

from langchain.text_splitter import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    separators=["\n\n", "\n", ". ", " ", ""],
    length_function=len,
)
chunks = splitter.split_text(document_text)

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("all-MiniLM-L6-v2")

def semantic_chunk(sentences: list[str], threshold: float = 0.75) -> list[str]:
    embeddings = model.encode(sentences)
    chunks = []
    current_chunk = [sentences[0]]

    for i in range(1, len(sentences)):
        similarity = np.dot(embeddings[i], embeddings[i-1]) / (
            np.linalg.norm(embeddings[i]) * np.linalg.norm(embeddings[i-1])
        )
        if similarity >= threshold:
            current_chunk.append(sentences[i])
        else:
            chunks.append(" ".join(current_chunk))
            current_chunk = [sentences[i]]

    if current_chunk:
        chunks.append(" ".join(current_chunk))
    return chunks

def structure_aware_chunk(elements: list, max_chunk_tokens: int = 500) -> list[dict]:
    chunks = []
    current_chunk = {"text": "", "metadata": {}, "tokens": 0}
    current_section = ""

    for element in elements:
        element_type = type(element).__name__
        element_text = element.text.strip()

        if not element_text:
            continue

        # Section headers create chunk boundaries
        if element_type in ("Title", "Header"):
            if current_chunk["text"]:
                chunks.append(current_chunk.copy())
            current_section = element_text
            current_chunk = {
                "text": element_text + "\n\n",
                "metadata": {"section": current_section},
                "tokens": count_tokens(element_text),
            }
            continue

        element_tokens = count_tokens(element_text)

        # Split if adding this element exceeds limit
        if current_chunk["tokens"] + element_tokens > max_chunk_tokens:
            if current_chunk["text"]:
                chunks.append(current_chunk.copy())
            current_chunk = {
                "text": element_text + "\n\n",
                "metadata": {"section": current_section},
                "tokens": element_tokens,
            }
        else:
            current_chunk["text"] += element_text + "\n\n"
            current_chunk["tokens"] += element_tokens

    if current_chunk["text"]:
        chunks.append(current_chunk)

    return chunks

# Deploying BGE-M3 on AWS with SageMaker
import sagemaker
from sagemaker.huggingface import HuggingFaceModel

hub_config = {
    "HF_MODEL_ID": "BAAI/bge-m3",
    "HF_TASK": "feature-extraction",
}

huggingface_model = HuggingFaceModel(
    transformers_version="4.37.0",
    pytorch_version="2.1.0",
    py_version="py310",
    env=hub_config,
    role=sagemaker.get_execution_role(),
)

predictor = huggingface_model.deploy(
    initial_instance_count=2,
    instance_type="ml.g5.xlarge",  # NVIDIA A10G, good price/performance
    endpoint_name="bge-m3-embeddings",
)

from openai import OpenAI

client = OpenAI()

response = client.embeddings.create(
    model="text-embedding-3-large",
    input="Cloud security best practices for multi-tenant environments",
    dimensions=1024,  # Reduced from 3072 default
)

# 1024 dimensions retains ~95% of retrieval quality at 1/3 storage cost
embedding = response.data[0].embedding

-- pgvector setup for RAG
CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE document_embeddings (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    content TEXT NOT NULL,
    embedding vector(1024) NOT NULL,
    metadata JSONB NOT NULL DEFAULT '{}',
    source_document TEXT NOT NULL,
    chunk_index INTEGER NOT NULL,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- HNSW index for approximate nearest neighbor search
CREATE INDEX ON document_embeddings
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 200);

-- Metadata index for filtered retrieval
CREATE INDEX idx_metadata ON document_embeddings USING GIN (metadata);

from qdrant_client import QdrantClient
from qdrant_client.models import (
    NamedSparseVector, NamedVector, Prefetch, Query, FusionQuery, Fusion
)

client = QdrantClient(url="http://qdrant:6333")

# Hybrid search with Reciprocal Rank Fusion (RRF)
results = client.query_points(
    collection_name="documents",
    prefetch=[
        # Dense vector search
        Prefetch(
            query=dense_embedding,
            using="dense",
            limit=20,
        ),
        # Sparse vector search (BM25)
        Prefetch(
            query=NamedSparseVector(
                name="bm25",
                vector=sparse_vector,
            ),
            using="bm25",
            limit=20,
        ),
    ],
    # Fuse results using RRF
    query=FusionQuery(fusion=Fusion.RRF),
    limit=10,
)

QUERY_REWRITE_PROMPT = """Given the user's question, generate 3 alternative
phrasings that might better match relevant documents. Focus on:
1. Technical terminology variations
2. More specific formulations
3. Broader context framing

User question: {query}

Return as a JSON array of strings."""

async def expand_query(query: str, llm_client) -> list[str]:
    response = await llm_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": QUERY_REWRITE_PROMPT.format(query=query)}],
        response_format={"type": "json_object"},
        temperature=0.3,
    )
    expanded = json.loads(response.choices[0].message.content)
    return [query] + expanded["queries"]  # Include original query

from sentence_transformers import CrossEncoder

reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")

def rerank_results(query: str, documents: list[dict], top_k: int = 5) -> list[dict]:
    pairs = [(query, doc["content"]) for doc in documents]
    scores = reranker.predict(pairs)

    for doc, score in zip(documents, scores):
        doc["rerank_score"] = float(score)

    ranked = sorted(documents, key=lambda x: x["rerank_score"], reverse=True)
    return ranked[:top_k]

def assemble_context(
    query: str,
    retrieved_chunks: list[dict],
    max_context_tokens: int = 6000,
    system_prompt_tokens: int = 500,
    max_response_tokens: int = 1500,
) -> str:
    # Reserve tokens for system prompt + query + response
    available_tokens = max_context_tokens - system_prompt_tokens - count_tokens(query) - max_response_tokens

    context_parts = []
    used_tokens = 0

    for chunk in retrieved_chunks:
        chunk_tokens = count_tokens(chunk["content"])
        if used_tokens + chunk_tokens > available_tokens:
            break

        source_label = f"[Source: {chunk['metadata'].get('source', 'Unknown')}, Section: {chunk['metadata'].get('section', 'N/A')}]"
        context_parts.append(f"{source_label}\n{chunk['content']}")
        used_tokens += chunk_tokens

    return "\n\n---\n\n".join(context_parts)

SYSTEM_PROMPT = """You are a technical assistant answering questions based on
the provided reference documents. Follow these rules strictly:

1. Answer ONLY based on the provided context. If the context does not contain
   sufficient information to answer, say "I don't have enough information to
   answer this question based on the available documents."
2. Cite your sources using [Source: filename, Section: section_name] format.
3. If multiple sources provide conflicting information, note the discrepancy.
4. Do not speculate or add information not present in the context.
5. Use precise technical terminology from the source documents."""

USER_TEMPLATE = """Context documents:
{context}

---

Question: {query}

Provide a detailed, accurate answer based on the above context."""

def evaluate_retrieval(
    test_queries: list[dict],  # {"query": str, "relevant_doc_ids": list[str]}
    retriever,
    k_values: list[int] = [1, 3, 5, 10],
) -> dict:
    metrics = {f"recall@{k}": [] for k in k_values}
    metrics.update({f"precision@{k}": [] for k in k_values})
    metrics["mrr"] = []

    for test_case in test_queries:
        results = retriever.retrieve(test_case["query"], top_k=max(k_values))
        retrieved_ids = [r["id"] for r in results]
        relevant_ids = set(test_case["relevant_doc_ids"])

        # MRR (Mean Reciprocal Rank)
        for rank, rid in enumerate(retrieved_ids, 1):
            if rid in relevant_ids:
                metrics["mrr"].append(1.0 / rank)
                break
        else:
            metrics["mrr"].append(0.0)

        # Recall@K and Precision@K
        for k in k_values:
            top_k_ids = set(retrieved_ids[:k])
            hits = len(top_k_ids & relevant_ids)
            metrics[f"recall@{k}"].append(hits / len(relevant_ids) if relevant_ids else 0)
            metrics[f"precision@{k}"].append(hits / k)

    return {key: sum(values) / len(values) for key, values in metrics.items()}

FAITHFULNESS_JUDGE_PROMPT = """Evaluate whether the ANSWER is faithfully
grounded in the provided CONTEXT.

CONTEXT:
{context}

ANSWER:
{answer}

Score on a scale of 1-5:
1: Answer contains fabricated information not in context
2: Answer mostly fabricated with minor grounding
3: Answer partially grounded, some claims unsupported
4: Answer mostly grounded, minor extrapolations
5: Answer fully grounded in provided context

Return JSON: {{"score": int, "reasoning": str, "unsupported_claims": list[str]}}"""

Monthly cost = Embedding costs + Vector storage costs + LLM generation costs

Example for 1M document corpus, 100K queries/month:

Embedding (one-time ingestion):
  1M docs × 500 tokens avg × $0.02/1M tokens = $10

Vector storage (monthly):
  1M vectors × 1024 dims × 4 bytes = 4 GB
  Pinecone: ~$70/month (s1.x1 pod)
  pgvector: ~$0 (existing PostgreSQL)
  Qdrant Cloud: ~$50/month

LLM generation (monthly):
  100K queries × 2000 input tokens × $2.50/1M tokens (GPT-4o) = $500
  100K queries × 500 output tokens × $10/1M tokens (GPT-4o) = $500

Total: ~$1,070/month with GPT-4o
       ~$120/month with GPT-4o-mini ($0.15/$0.60 per 1M tokens)