from qdrant_client import QdrantClient, models
from openai import OpenAI

qdrant = QdrantClient(url="http://localhost:6333")
openai = OpenAI()

def embed(text: str) -> list[float]:
    resp = openai.embeddings.create(model="text-embedding-3-large", input=text)
    return resp.data[0].embedding

def hybrid_search(query: str, collection: str, top_k: int = 10) -> list:
    # 1. Dense (semantic) search — finds conceptually similar content
    dense_results = qdrant.search(
        collection_name=collection,
        query_vector=embed(query),
        limit=top_k,
        with_payload=True
    )

    # 2. Sparse (keyword) search — finds exact term matches
    sparse_results = qdrant.search(
        collection_name=collection,
        query=query,          # Full-text BM25 search
        limit=top_k,
        with_payload=True
    )

    # 3. Reciprocal Rank Fusion — merge both result sets
    return rrf_merge(dense_results, sparse_results, k=60)

def rrf_merge(list1, list2, k: int = 60) -> list:
    """Reciprocal Rank Fusion — standard hybrid search merger"""
    scores = {}
    for rank, doc in enumerate(list1):
        scores[doc.id] = scores.get(doc.id, 0) + 1 / (k + rank + 1)
    for rank, doc in enumerate(list2):
        scores[doc.id] = scores.get(doc.id, 0) + 1 / (k + rank + 1)
    return sorted(scores.items(), key=lambda x: x[1], reverse=True)

import cohere

co = cohere.Client(api_key="...")

def rerank_and_trim(
    query: str,
    documents: list[str],
    top_n: int = 5
) -> list[str]:
    """
    Take 20 retrieved chunks, return the 5 most relevant.
    Cross-encoder re-ranking often improves precision.
    Verify current provider pricing before production use.
    Quality improvement is usually strongest on complex or ambiguous queries.
    """
    results = co.rerank(
        model="rerank-english-v3.0",
        query=query,
        documents=documents,
        top_n=top_n,
        return_documents=True
    )
    return [r.document.text for r in results.results]

# Full pipeline:
raw_chunks   = hybrid_search(query, collection="codebase", top_k=20)
ranked_chunks = rerank_and_trim(query, raw_chunks, top_n=5)
answer       = llm.complete(f"Answer using these docs:\n{ranked_chunks}\n\nQ: {query}")

from ragas import evaluate
from ragas.metrics import faithfulness, answer_relevancy, context_precision, context_recall
from datasets import Dataset

# Build evaluation dataset
eval_data = {
    "question": ["Where is rate limiting implemented?", ...],
    "answer":   [agent_answer, ...],           # What the agent said
    "contexts": [[chunk1, chunk2, ...], ...],   # What was retrieved
    "ground_truth": ["In api/middleware/ratelimit.py", ...] # Correct answer
}
dataset = Dataset.from_dict(eval_data)

# Run RAGAS evaluation
result = evaluate(
    dataset=dataset,
    metrics=[faithfulness, answer_relevancy, context_precision, context_recall]
)
print(result)
# {'faithfulness': 0.89, 'answer_relevancy': 0.84, ...}

import numpy as np

def cosine_similarity(v1: list, v2: list) -> float:
    """Standard measure of vector similarity for embeddings."""
    a, b = np.array(v1), np.array(v2)
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# Vector DBs store cosine DISTANCE (1 - similarity) for sorting
# distance=0.0 → identical, distance=1.0 → completely unrelated

# text-embedding-3-large natively supports dimension reduction
resp = openai.embeddings.create(
    model="text-embedding-3-large",
    input=text,
    dimensions=256   # Truncate to 256 dims instead of default 3072
               # Smaller storage and faster search
               # Measure quality loss on your own eval set
)
# Great for: high-volume applications where cost and latency matter
# Avoid for: high-stakes retrieval where every percentage matters

from qdrant_client import QdrantClient, models
import uuid

client = QdrantClient(
    url="https://your-cluster.qdrant.io",  # Or "http://localhost:6333"
    api_key="your-api-key"
)

# 1. Create collection with HNSW index
client.create_collection(
    collection_name="documents",
    vectors_config=models.VectorParams(
        size=1536,                          # Match your embedding model
        distance=models.Distance.COSINE,     # Standard for text embeddings
        hnsw_config=models.HnswConfigDiff(
            m=16,                             # More connections = better recall, more RAM
            ef_construct=100                 # Higher = better index quality, slower build
        )
    ),
    # Payload index for fast metadata filtering
    on_disk_payload=True                     # Store metadata on disk, vectors in RAM
)

# 2. Create payload index for fast filtering
client.create_payload_index(
    collection_name="documents",
    field_name="user_id",
    field_schema=models.PayloadSchemaType.KEYWORD
)

# 3. Upsert vectors with rich metadata (payload)
client.upsert(
    collection_name="documents",
    points=[
        models.PointStruct(
            id=str(uuid.uuid4()),
            vector=embed(chunk_text),        # 1536-dim float list
            payload={
                "text": chunk_text,
                "source": "confluence.md",
                "user_id": "prem",
                "section": "architecture",
                "date": "2026-01-15",
                "chunk_index": 3
            }
        )
        for chunk_text in chunks
    ]
)

# 4. Search with metadata filtering (ALWAYS filter before vector search!)
results = client.search(
    collection_name="documents",
    query_vector=embed("Where is rate limiting implemented?"),
    query_filter=models.Filter(
        must=[
            models.FieldCondition(key="user_id", match=models.MatchValue(value="prem")),
            models.FieldCondition(key="section", match=models.MatchValue(value="architecture"))
        ]
    ),
    limit=10,
    with_payload=True,
    search_params=models.SearchParams(ef=128)  # Higher ef = better recall at query time
)

for result in results:
    print(f"Score: {result.score:.3f} | {result.payload['text'][:100]}")

-- Install extension (once per database)
CREATE EXTENSION IF NOT EXISTS vector;

-- Create table with vector column
CREATE TABLE documents (
    id          UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    content     TEXT NOT NULL,
    source      TEXT,
    section     TEXT,
    user_id     TEXT,
    created_at  TIMESTAMPTZ DEFAULT NOW(),
    embedding   vector(1536)     -- Dimension must match your model
);

-- HNSW index for fast approximate search (Postgres 16+)
CREATE INDEX ON documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- Standard indexes for metadata filtering
CREATE INDEX ON documents (user_id);
CREATE INDEX ON documents (section);

-- Semantic search with metadata filter
-- The $1 parameter is your query embedding as a vector literal
SELECT
    content,
    source,
    1 - (embedding <=> $1) AS similarity   -- <=> is cosine distance
FROM documents
WHERE
    user_id = $2                              -- Filter BEFORE vector search
    AND section = $3
    AND created_at > NOW() - INTERVAL '30 days'
ORDER BY embedding <=> $1                    -- Sort by similarity
LIMIT 10;

-- Full-text + vector hybrid search
SELECT content, source,
    1 - (embedding <=> $1) AS vec_score,
    ts_rank(to_tsvector('english', content), plainto_tsquery($2)) AS text_score
FROM documents
WHERE to_tsvector('english', content) @@ plainto_tsquery($2)
ORDER BY (vec_score * 0.7) + (text_score * 0.3) DESC   -- Weighted hybrid
LIMIT 10;

from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.graph import StateGraph

# PostgresSaver persists every step to Postgres
# Enables: pause/resume, HITL, time-travel debugging, rollback
checkpointer = PostgresSaver.from_conn_string(
    "postgresql://user:pass@localhost/agentdb"
)

# Compile graph with checkpointer
graph = build_agent_graph()
app = graph.compile(checkpointer=checkpointer)

# Run with thread_id — state persists across calls
config = {"configurable": {"thread_id": "prem-session-2026-07"}}

# First message
result1 = app.invoke(
    {"messages": [("user", "Research Adyen for me")]},
    config=config
)

# Hours later — agent remembers the entire prior conversation
result2 = app.invoke(
    {"messages": [("user", "Now compare them with Stripe")]},
    config=config  # Same thread_id → same state is loaded
)

# Time-travel: inspect state at any past step
history = list(app.get_state_history(config))
past_state = app.get_state(config, checkpoint_id=history[3].checkpoint_id)

# Available checkpointers:
# MemorySaver     — in-memory, lost on restart (dev only)
# SqliteSaver     — SQLite file (single-process)
# PostgresSaver   — PostgreSQL-backed durable state
# RedisSaver      — Redis (low-latency, high-throughput)

from mem0 import Memory
from mem0 import MemoryClient  # Cloud managed version

# Initialise (uses Qdrant + OpenAI embeddings by default)
m = Memory()

USER_ID = "prem"

# ── STORING MEMORIES ─────────────────────────────────────────

# From a conversation (Mem0 extracts facts automatically)
m.add(
    messages=[
        {"role": "user",      "content": "I prefer TypeScript over Python for backends"},
        {"role": "assistant", "content": "Got it, I'll use TypeScript for your projects"}
    ],
    user_id=USER_ID
)
# Mem0 extracts: "User prefers TypeScript over Python for backend development"
# and stores it with a vector embedding

# From an explicit fact
m.add(
    messages=[{"role": "user", "content": "I decided to use LangGraph for the PR review pipeline"}],
    user_id=USER_ID,
    metadata={"project": "pr-agent", "decision": True}
)

# ── RETRIEVING MEMORIES ──────────────────────────────────────

# Semantic search — finds relevant memories
memories = m.search(
    query="What programming language should I use?",
    user_id=USER_ID,
    limit=5
)
# Returns: [{"memory": "User prefers TypeScript over Python", "score": 0.94}]

# Get all memories for a user
all_memories = m.get_all(user_id=USER_ID)

# ── MEMORY CRUD ──────────────────────────────────────────────

# Delete a specific memory (user privacy control)
m.delete(memory_id="mem_abc123")

# Delete all memories for a user (GDPR deletion)
m.delete_all(user_id=USER_ID)

# Update a memory when facts change
m.update(
    memory_id="mem_abc123",
    data="User now prefers Python for data pipelines, TypeScript for APIs"
)

from datetime import datetime, timedelta
import json

class EpisodicMemoryStore:
    """
    Store and retrieve past agent episodes.
    Each episode = one completed agent run with its outcome.
    """

    def __init__(self, qdrant_client, embed_fn):
        self.qdrant = qdrant_client
        self.embed  = embed_fn
        self.collection = "episodic_memory"

    def store_episode(
        self,
        user_id: str,
        query: str,
        outcome: str,
        tools_used: list[str],
        result_summary: str
    ):
        """Store a completed agent run as an episode."""
        episode_text = f"Task: {query}\nOutcome: {outcome}\nSummary: {result_summary}"
        self.qdrant.upsert(
            collection_name=self.collection,
            points=[models.PointStruct(
                id=str(uuid.uuid4()),
                vector=self.embed(episode_text),
                payload={
                    "user_id":     user_id,
                    "query":       query,
                    "outcome":     outcome,
                    "tools_used":  tools_used,
                    "summary":     result_summary,
                    "timestamp":   datetime.utcnow().isoformat()
                }
            )]
        )

    def recall_relevant(self, user_id: str, current_query: str, limit: int = 3) -> list:
        """Find past episodes relevant to the current task."""
        results = self.qdrant.search(
            collection_name=self.collection,
            query_vector=self.embed(current_query),
            query_filter=models.Filter(must=[
                models.FieldCondition(key="user_id", match=models.MatchValue(value=user_id)),
                # Only recall recent episodes (configurable decay)
                models.FieldCondition(
                    key="timestamp",
                    range=models.DatetimeRange(
                        gte=datetime.utcnow() - timedelta(days=90)
                    )
                )
            ]),
            limit=limit
        )
        return [r.payload for r in results]

    def inject_into_prompt(self, user_id: str, current_query: str) -> str:
        """Build memory context string to prepend to agent system prompt."""
        episodes = self.recall_relevant(user_id, current_query)
        if not episodes:
            return ""
        memory_text = "\n\n## Relevant past work:\n"
        for ep in episodes:
            memory_text += f"\n- [{ep['timestamp'][:10]}] {ep['query']}: {ep['summary']}"
        return memory_text