AST and Symbol-Aware Retrieval (Tier 2)

# retrieval/embed_ast_chunks.py
"""Embed AST-aware chunks and store in a separate Qdrant collection."""
import json
from pathlib import Path
from openai import OpenAI
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

CHUNKS_PATH = Path("retrieval/chunks_ast.jsonl")
COLLECTION_NAME = "anchor-repo-ast"
EMBEDDING_MODEL = "text-embedding-3-small"
EMBEDDING_DIM = 1536
BATCH_SIZE = 50

client = OpenAI()
qdrant = QdrantClient(path="retrieval/qdrant_data")


def load_chunks():
    """Load AST-aware chunks from the JSONL chunk store.

    Args:
        None.

    Returns:
        A list of chunk dictionaries from the chunk store.
    """
    chunks = []
    with open(CHUNKS_PATH) as f:
        for line in f:
            if line.strip():
                chunks.append(json.loads(line))
    return chunks


def embed_texts(texts):
    """Generate embeddings for a batch of AST-aware chunk texts.

    Args:
        texts: Chunk texts to embed.

    Returns:
        A list of embedding vectors in the same order as the input texts.
    """
    response = client.embeddings.create(model=EMBEDDING_MODEL, input=texts)
    return [item.embedding for item in response.data]


def create_and_store():
    """Embed AST-aware chunks and store them in the Qdrant collection.

    Args:
        None.

    Returns:
        None. The vector collection is recreated and populated in place.
    """
    chunks = load_chunks()
    collections = [c.name for c in qdrant.get_collections().collections]
    if COLLECTION_NAME in collections:
        qdrant.delete_collection(COLLECTION_NAME)
    qdrant.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE),
    )
    print(f"Created collection '{COLLECTION_NAME}'")

    for batch_start in range(0, len(chunks), BATCH_SIZE):
        batch = chunks[batch_start:batch_start + BATCH_SIZE]
        texts = [f"{c['symbol_type']} {c['symbol_name']} in {c['file_path']}\n\n{c['text']}" for c in batch]
        embeddings = embed_texts(texts)
        points = [
            PointStruct(id=batch_start + i, vector=emb, payload={
                "chunk_id": chunk["chunk_id"], "file_path": chunk["file_path"],
                "symbol_name": chunk["symbol_name"], "symbol_type": chunk["symbol_type"],
                "start_line": chunk["start_line"], "end_line": chunk["end_line"],
                "text": chunk["text"],
            })
            for i, (chunk, emb) in enumerate(zip(batch, embeddings))
        ]
        qdrant.upsert(collection_name=COLLECTION_NAME, points=points)
        print(f"  Stored {batch_start + len(batch)}/{len(chunks)} chunks")
    print(f"\nDone. {len(chunks)} AST-aware chunks stored in '{COLLECTION_NAME}'")


if __name__ == "__main__":
    create_and_store()

# retrieval/embed_ast_chunks.py
"""Embed AST-aware chunks and store in a separate Qdrant collection."""
import json
from pathlib import Path
from google import genai
from google.genai import types
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

CHUNKS_PATH = Path("retrieval/chunks_ast.jsonl")
COLLECTION_NAME = "anchor-repo-ast"
EMBEDDING_MODEL = "gemini-embedding-001"
EMBEDDING_DIM = 768
BATCH_SIZE = 50

client = genai.Client()
qdrant = QdrantClient(path="retrieval/qdrant_data")


def load_chunks():
    """Load AST-aware chunks from the JSONL chunk store.

    Args:
        None.

    Returns:
        A list of chunk dictionaries from the chunk store.
    """
    chunks = []
    with open(CHUNKS_PATH) as f:
        for line in f:
            if line.strip():
                chunks.append(json.loads(line))
    return chunks


def embed_texts(texts):
    """Generate embeddings for a batch of AST-aware chunk texts.

    Args:
        texts: Chunk texts to embed.

    Returns:
        A list of embedding vectors in the same order as the input texts.
    """
    response = client.models.embed_content(
        model=EMBEDDING_MODEL,
        contents=texts,
        config=types.EmbedContentConfig(
            task_type="RETRIEVAL_DOCUMENT",
            output_dimensionality=EMBEDDING_DIM,
        ),
    )
    return [item.values for item in response.embeddings]


def create_and_store():
    """Embed AST-aware chunks and store them in the Qdrant collection.

    Args:
        None.

    Returns:
        None. The vector collection is recreated and populated in place.
    """
    chunks = load_chunks()
    collections = [c.name for c in qdrant.get_collections().collections]
    if COLLECTION_NAME in collections:
        qdrant.delete_collection(COLLECTION_NAME)
    qdrant.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE),
    )
    print(f"Created collection '{COLLECTION_NAME}'")

    for batch_start in range(0, len(chunks), BATCH_SIZE):
        batch = chunks[batch_start:batch_start + BATCH_SIZE]
        texts = [f"{c['symbol_type']} {c['symbol_name']} in {c['file_path']}\n\n{c['text']}" for c in batch]
        embeddings = embed_texts(texts)
        points = [
            PointStruct(id=batch_start + i, vector=emb, payload={
                "chunk_id": chunk["chunk_id"], "file_path": chunk["file_path"],
                "symbol_name": chunk["symbol_name"], "symbol_type": chunk["symbol_type"],
                "start_line": chunk["start_line"], "end_line": chunk["end_line"],
                "text": chunk["text"],
            })
            for i, (chunk, emb) in enumerate(zip(batch, embeddings))
        ]
        qdrant.upsert(collection_name=COLLECTION_NAME, points=points)
        print(f"  Stored {batch_start + len(batch)}/{len(chunks)} chunks")
    print(f"\nDone. {len(chunks)} AST-aware chunks stored in '{COLLECTION_NAME}'")


if __name__ == "__main__":
    create_and_store()

# retrieval/embed_ast_chunks.py
"""Embed AST-aware chunks.

Anthropic does not provide an embeddings model, so this script
uses OpenAI for embeddings. This mixed-provider pattern is normal
in AI engineering. Use Anthropic for generation in other scripts.
"""
import json
from pathlib import Path
from openai import OpenAI
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

CHUNKS_PATH = Path("retrieval/chunks_ast.jsonl")
COLLECTION_NAME = "anchor-repo-ast"
EMBEDDING_MODEL = "text-embedding-3-small"
EMBEDDING_DIM = 1536
BATCH_SIZE = 50

# Anthropic has no embeddings API. Use OpenAI for embeddings here.
client = OpenAI()
qdrant = QdrantClient(path="retrieval/qdrant_data")


def load_chunks():
    """Load AST-aware chunks from the JSONL chunk store.

    Args:
        None.

    Returns:
        A list of chunk dictionaries from the chunk store.
    """
    chunks = []
    with open(CHUNKS_PATH) as f:
        for line in f:
            if line.strip():
                chunks.append(json.loads(line))
    return chunks


def embed_texts(texts):
    """Generate embeddings for a batch of AST-aware chunk texts.

    Args:
        texts: Chunk texts to embed.

    Returns:
        A list of embedding vectors in the same order as the input texts.
    """
    response = client.embeddings.create(model=EMBEDDING_MODEL, input=texts)
    return [item.embedding for item in response.data]


def create_and_store():
    """Embed AST-aware chunks and store them in the Qdrant collection.

    Args:
        None.

    Returns:
        None. The vector collection is recreated and populated in place.
    """
    chunks = load_chunks()
    collections = [c.name for c in qdrant.get_collections().collections]
    if COLLECTION_NAME in collections:
        qdrant.delete_collection(COLLECTION_NAME)
    qdrant.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE),
    )
    print(f"Created collection '{COLLECTION_NAME}'")

    for batch_start in range(0, len(chunks), BATCH_SIZE):
        batch = chunks[batch_start:batch_start + BATCH_SIZE]
        texts = [f"{c['symbol_type']} {c['symbol_name']} in {c['file_path']}\n\n{c['text']}" for c in batch]
        embeddings = embed_texts(texts)
        points = [
            PointStruct(id=batch_start + i, vector=emb, payload={
                "chunk_id": chunk["chunk_id"], "file_path": chunk["file_path"],
                "symbol_name": chunk["symbol_name"], "symbol_type": chunk["symbol_type"],
                "start_line": chunk["start_line"], "end_line": chunk["end_line"],
                "text": chunk["text"],
            })
            for i, (chunk, emb) in enumerate(zip(batch, embeddings))
        ]
        qdrant.upsert(collection_name=COLLECTION_NAME, points=points)
        print(f"  Stored {batch_start + len(batch)}/{len(chunks)} chunks")
    print(f"\nDone. {len(chunks)} AST-aware chunks stored in '{COLLECTION_NAME}'")


if __name__ == "__main__":
    create_and_store()

# retrieval/embed_ast_chunks.py
"""Embed AST-aware chunks using Hugging Face Inference Providers."""
import json
import os
from pathlib import Path
from huggingface_hub import InferenceClient
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

CHUNKS_PATH = Path("retrieval/chunks_ast.jsonl")
COLLECTION_NAME = "anchor-repo-ast"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
EMBEDDING_DIM = 384
BATCH_SIZE = 50

client = InferenceClient(token=os.environ["HF_TOKEN"])
qdrant = QdrantClient(path="retrieval/qdrant_data")


def load_chunks():
    """Load AST-aware chunks from the JSONL chunk store.

    Args:
        None.

    Returns:
        A list of chunk dictionaries from the chunk store.
    """
    chunks = []
    with open(CHUNKS_PATH) as f:
        for line in f:
            if line.strip():
                chunks.append(json.loads(line))
    return chunks


def embed_texts(texts):
    """Generate embeddings for a batch of AST-aware chunk texts.

    Args:
        texts: Chunk texts to embed.

    Returns:
        A list of embedding vectors in the same order as the input texts.
    """
    embeddings = []
    for text in texts:
        vec = client.feature_extraction(text, model=EMBEDDING_MODEL)
        embeddings.append(vec.tolist() if hasattr(vec, "tolist") else vec)
    return embeddings


def create_and_store():
    """Embed AST-aware chunks and store them in the Qdrant collection.

    Args:
        None.

    Returns:
        None. The vector collection is recreated and populated in place.
    """
    chunks = load_chunks()
    collections = [c.name for c in qdrant.get_collections().collections]
    if COLLECTION_NAME in collections:
        qdrant.delete_collection(COLLECTION_NAME)
    qdrant.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE),
    )
    print(f"Created collection '{COLLECTION_NAME}'")

    for batch_start in range(0, len(chunks), BATCH_SIZE):
        batch = chunks[batch_start:batch_start + BATCH_SIZE]
        texts = [f"{c['symbol_type']} {c['symbol_name']} in {c['file_path']}\n\n{c['text']}" for c in batch]
        embeddings = embed_texts(texts)
        points = [
            PointStruct(id=batch_start + i, vector=emb, payload={
                "chunk_id": chunk["chunk_id"], "file_path": chunk["file_path"],
                "symbol_name": chunk["symbol_name"], "symbol_type": chunk["symbol_type"],
                "start_line": chunk["start_line"], "end_line": chunk["end_line"],
                "text": chunk["text"],
            })
            for i, (chunk, emb) in enumerate(zip(batch, embeddings))
        ]
        qdrant.upsert(collection_name=COLLECTION_NAME, points=points)
        print(f"  Stored {batch_start + len(batch)}/{len(chunks)} chunks")
    print(f"\nDone. {len(chunks)} AST-aware chunks stored in '{COLLECTION_NAME}'")


if __name__ == "__main__":
    create_and_store()

import os
from openai import OpenAI

# Same as the OpenAI embed_ast_chunks.py file,
# but point the embedding client at GitHub Models.
EMBEDDING_MODEL = "openai/text-embedding-3-small"
EMBEDDING_DIM = 1536

client = OpenAI(
    base_url="https://models.github.ai/inference",
    api_key=os.environ["GITHUB_TOKEN"],
    default_headers={
        "Accept": "application/vnd.github+json",
        "X-GitHub-Api-Version": "2026-03-10",
    },
)
# Keep the rest of the OpenAI file unchanged:
# same AST chunks, same Qdrant structure, same batched upserts.

# retrieval/embed_ast_chunks.py
"""Embed AST-aware chunks using Ollama."""
import json
from pathlib import Path
from ollama import Client
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

CHUNKS_PATH = Path("retrieval/chunks_ast.jsonl")
COLLECTION_NAME = "anchor-repo-ast"
EMBEDDING_MODEL = "nomic-embed-text:latest"
EMBEDDING_DIM = 768
BATCH_SIZE = 50

# Run `ollama pull nomic-embed-text:latest` first.
client = Client(host="http://localhost:11434")
qdrant = QdrantClient(path="retrieval/qdrant_data")


def load_chunks():
    """Load AST-aware chunks from the JSONL chunk store.

    Args:
        None.

    Returns:
        A list of chunk dictionaries from the chunk store.
    """
    chunks = []
    with open(CHUNKS_PATH) as f:
        for line in f:
            if line.strip():
                chunks.append(json.loads(line))
    return chunks


def embed_texts(texts):
    """Generate embeddings for a batch of AST-aware chunk texts.

    Args:
        texts: Chunk texts to embed.

    Returns:
        A list of embedding vectors in the same order as the input texts.
    """
    embeddings = []
    for text in texts:
        response = client.embed(model=EMBEDDING_MODEL, input=text)
        embeddings.append(response.embeddings[0])
    return embeddings


def create_and_store():
    """Embed AST-aware chunks and store them in the Qdrant collection.

    Args:
        None.

    Returns:
        None. The vector collection is recreated and populated in place.
    """
    chunks = load_chunks()
    collections = [c.name for c in qdrant.get_collections().collections]
    if COLLECTION_NAME in collections:
        qdrant.delete_collection(COLLECTION_NAME)
    qdrant.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE),
    )
    print(f"Created collection '{COLLECTION_NAME}'")

    for batch_start in range(0, len(chunks), BATCH_SIZE):
        batch = chunks[batch_start:batch_start + BATCH_SIZE]
        texts = [f"{c['symbol_type']} {c['symbol_name']} in {c['file_path']}\n\n{c['text']}" for c in batch]
        embeddings = embed_texts(texts)
        points = [
            PointStruct(id=batch_start + i, vector=emb, payload={
                "chunk_id": chunk["chunk_id"], "file_path": chunk["file_path"],
                "symbol_name": chunk["symbol_name"], "symbol_type": chunk["symbol_type"],
                "start_line": chunk["start_line"], "end_line": chunk["end_line"],
                "text": chunk["text"],
            })
            for i, (chunk, emb) in enumerate(zip(batch, embeddings))
        ]
        qdrant.upsert(collection_name=COLLECTION_NAME, points=points)
        print(f"  Stored {batch_start + len(batch)}/{len(chunks)} chunks")
    print(f"\nDone. {len(chunks)} AST-aware chunks stored in '{COLLECTION_NAME}'")


if __name__ == "__main__":
    create_and_store()

# retrieval/compare_tiers.py
"""Compare naive vs. AST-aware retrieval on benchmark questions."""
import json
from pathlib import Path
from openai import OpenAI
from qdrant_client import QdrantClient

BENCHMARK_FILE = Path("benchmark-questions.jsonl")
EMBEDDING_MODEL = "text-embedding-3-small"
TOP_K = 5

client = OpenAI()
qdrant = QdrantClient(path="retrieval/qdrant_data")


def retrieve_from(collection, query, top_k=TOP_K):
    """Query one vector collection and return the top AST-aware matches.

    Args:
        collection: Qdrant collection name to query.
        query: User question or lookup string to embed.
        top_k: Number of matches to return.

    Returns:
        A list of ranked retrieval hits with file, symbol, and preview metadata.
    """
    response = client.embeddings.create(model=EMBEDDING_MODEL, input=[query])
    query_vector = response.data[0].embedding
    results = qdrant.query_points(collection_name=collection, query=query_vector, limit=top_k)
    return [{"file_path": hit.payload["file_path"], "score": round(hit.score, 4),
             "text_preview": hit.payload["text"][:120],
             "symbol_name": hit.payload.get("symbol_name", "n/a")}
            for hit in results.points]


def compare():
    """Compare naive retrieval against AST-aware retrieval on benchmark questions.

    Args:
        None.

    Returns:
        None. Results are printed for manual inspection.
    """
    questions = []
    with open(BENCHMARK_FILE) as f:
        for line in f:
            if line.strip():
                questions.append(json.loads(line))

    print(f"Comparing naive vs. AST-aware retrieval on {min(len(questions), 15)} questions\n")
    for q in questions[:15]:
        print(f"[{q['category']}] {q['question'][:70]}")
        naive = retrieve_from("anchor-repo-naive", q["question"])
        ast_aware = retrieve_from("anchor-repo-ast", q["question"])
        print(f"  Naive:     {', '.join(set(r['file_path'] for r in naive))}")
        print(f"  AST-aware: {', '.join(set(r['file_path'] for r in ast_aware))}")
        for r in ast_aware[:3]:
            print(f"    [{r['score']}] {r['symbol_name']} in {r['file_path']}")
        print()


if __name__ == "__main__":
    compare()

# retrieval/compare_tiers.py
"""Compare naive vs. AST-aware retrieval on benchmark questions."""
import json
from pathlib import Path
from google import genai
from google.genai import types
from qdrant_client import QdrantClient

BENCHMARK_FILE = Path("benchmark-questions.jsonl")
EMBEDDING_MODEL = "gemini-embedding-001"
EMBEDDING_DIM = 768
TOP_K = 5

client = genai.Client()
qdrant = QdrantClient(path="retrieval/qdrant_data")


def retrieve_from(collection, query, top_k=TOP_K):
    """Query one vector collection and return the top AST-aware matches.

    Args:
        collection: Qdrant collection name to query.
        query: User question or lookup string to embed.
        top_k: Number of matches to return.

    Returns:
        A list of ranked retrieval hits with file, symbol, and preview metadata.
    """
    response = client.models.embed_content(
        model=EMBEDDING_MODEL,
        contents=[query],
        config=types.EmbedContentConfig(
            task_type="RETRIEVAL_QUERY",
            output_dimensionality=EMBEDDING_DIM,
        ),
    )
    query_vector = response.embeddings[0].values
    results = qdrant.query_points(collection_name=collection, query=query_vector, limit=top_k)
    return [{"file_path": hit.payload["file_path"], "score": round(hit.score, 4),
             "text_preview": hit.payload["text"][:120],
             "symbol_name": hit.payload.get("symbol_name", "n/a")}
            for hit in results.points]


def compare():
    """Compare naive retrieval against AST-aware retrieval on benchmark questions.

    Args:
        None.

    Returns:
        None. Results are printed for manual inspection.
    """
    questions = []
    with open(BENCHMARK_FILE) as f:
        for line in f:
            if line.strip():
                questions.append(json.loads(line))

    print(f"Comparing naive vs. AST-aware retrieval on {min(len(questions), 15)} questions\n")
    for q in questions[:15]:
        print(f"[{q['category']}] {q['question'][:70]}")
        naive = retrieve_from("anchor-repo-naive", q["question"])
        ast_aware = retrieve_from("anchor-repo-ast", q["question"])
        print(f"  Naive:     {', '.join(set(r['file_path'] for r in naive))}")
        print(f"  AST-aware: {', '.join(set(r['file_path'] for r in ast_aware))}")
        for r in ast_aware[:3]:
            print(f"    [{r['score']}] {r['symbol_name']} in {r['file_path']}")
        print()


if __name__ == "__main__":
    compare()

# retrieval/compare_tiers.py
"""Compare naive vs. AST-aware retrieval.

Anthropic does not provide an embeddings model, so this script
uses OpenAI for the embedding query. This is the same mixed-provider
pattern from embed_ast_chunks.py.
"""
import json
from pathlib import Path
from openai import OpenAI
from qdrant_client import QdrantClient

BENCHMARK_FILE = Path("benchmark-questions.jsonl")
EMBEDDING_MODEL = "text-embedding-3-small"
TOP_K = 5

# Anthropic has no embeddings API. Use OpenAI for embeddings here.
client = OpenAI()
qdrant = QdrantClient(path="retrieval/qdrant_data")


def retrieve_from(collection, query, top_k=TOP_K):
    """Query one vector collection and return the top AST-aware matches.

    Args:
        collection: Qdrant collection name to query.
        query: User question or lookup string to embed.
        top_k: Number of matches to return.

    Returns:
        A list of ranked retrieval hits with file, symbol, and preview metadata.
    """
    response = client.embeddings.create(model=EMBEDDING_MODEL, input=[query])
    query_vector = response.data[0].embedding
    results = qdrant.query_points(collection_name=collection, query=query_vector, limit=top_k)
    return [{"file_path": hit.payload["file_path"], "score": round(hit.score, 4),
             "text_preview": hit.payload["text"][:120],
             "symbol_name": hit.payload.get("symbol_name", "n/a")}
            for hit in results.points]


def compare():
    """Compare naive retrieval against AST-aware retrieval on benchmark questions.

    Args:
        None.

    Returns:
        None. Results are printed for manual inspection.
    """
    questions = []
    with open(BENCHMARK_FILE) as f:
        for line in f:
            if line.strip():
                questions.append(json.loads(line))

    print(f"Comparing naive vs. AST-aware retrieval on {min(len(questions), 15)} questions\n")
    for q in questions[:15]:
        print(f"[{q['category']}] {q['question'][:70]}")
        naive = retrieve_from("anchor-repo-naive", q["question"])
        ast_aware = retrieve_from("anchor-repo-ast", q["question"])
        print(f"  Naive:     {', '.join(set(r['file_path'] for r in naive))}")
        print(f"  AST-aware: {', '.join(set(r['file_path'] for r in ast_aware))}")
        for r in ast_aware[:3]:
            print(f"    [{r['score']}] {r['symbol_name']} in {r['file_path']}")
        print()


if __name__ == "__main__":
    compare()

# retrieval/compare_tiers.py
"""Compare naive vs. AST-aware retrieval using Hugging Face embeddings."""
import json
import os
from pathlib import Path
from huggingface_hub import InferenceClient
from qdrant_client import QdrantClient

BENCHMARK_FILE = Path("benchmark-questions.jsonl")
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
TOP_K = 5

client = InferenceClient(token=os.environ["HF_TOKEN"])
qdrant = QdrantClient(path="retrieval/qdrant_data")


def retrieve_from(collection, query, top_k=TOP_K):
    """Query one vector collection and return the top AST-aware matches.

    Args:
        collection: Qdrant collection name to query.
        query: User question or lookup string to embed.
        top_k: Number of matches to return.

    Returns:
        A list of ranked retrieval hits with file, symbol, and preview metadata.
    """
    vec = client.feature_extraction(query, model=EMBEDDING_MODEL)
    query_vector = vec.tolist() if hasattr(vec, "tolist") else vec
    results = qdrant.query_points(collection_name=collection, query=query_vector, limit=top_k)
    return [{"file_path": hit.payload["file_path"], "score": round(hit.score, 4),
             "text_preview": hit.payload["text"][:120],
             "symbol_name": hit.payload.get("symbol_name", "n/a")}
            for hit in results.points]


def compare():
    """Compare naive retrieval against AST-aware retrieval on benchmark questions.

    Args:
        None.

    Returns:
        None. Results are printed for manual inspection.
    """
    questions = []
    with open(BENCHMARK_FILE) as f:
        for line in f:
            if line.strip():
                questions.append(json.loads(line))

    print(f"Comparing naive vs. AST-aware retrieval on {min(len(questions), 15)} questions\n")
    for q in questions[:15]:
        print(f"[{q['category']}] {q['question'][:70]}")
        naive = retrieve_from("anchor-repo-naive", q["question"])
        ast_aware = retrieve_from("anchor-repo-ast", q["question"])
        print(f"  Naive:     {', '.join(set(r['file_path'] for r in naive))}")
        print(f"  AST-aware: {', '.join(set(r['file_path'] for r in ast_aware))}")
        for r in ast_aware[:3]:
            print(f"    [{r['score']}] {r['symbol_name']} in {r['file_path']}")
        print()


if __name__ == "__main__":
    compare()

import os
from openai import OpenAI

# Same as the OpenAI compare_tiers.py file,
# but point the embedding client at GitHub Models.
EMBEDDING_MODEL = "openai/text-embedding-3-small"

client = OpenAI(
    base_url="https://models.github.ai/inference",
    api_key=os.environ["GITHUB_TOKEN"],
    default_headers={
        "Accept": "application/vnd.github+json",
        "X-GitHub-Api-Version": "2026-03-10",
    },
)
# Keep the rest of the comparison logic unchanged:
# query each collection the same way, print the top hits,
# and compare how AST-aware retrieval improves over the naive baseline.

# retrieval/compare_tiers.py
"""Compare naive vs. AST-aware retrieval using Ollama embeddings."""
import json
from pathlib import Path
from ollama import Client
from qdrant_client import QdrantClient

BENCHMARK_FILE = Path("benchmark-questions.jsonl")
EMBEDDING_MODEL = "nomic-embed-text:latest"
TOP_K = 5

# Run `ollama pull nomic-embed-text:latest` first.
client = Client(host="http://localhost:11434")
qdrant = QdrantClient(path="retrieval/qdrant_data")


def retrieve_from(collection, query, top_k=TOP_K):
    """Query one vector collection and return the top AST-aware matches.

    Args:
        collection: Qdrant collection name to query.
        query: User question or lookup string to embed.
        top_k: Number of matches to return.

    Returns:
        A list of ranked retrieval hits with file, symbol, and preview metadata.
    """
    response = client.embed(model=EMBEDDING_MODEL, input=query)
    query_vector = response.embeddings[0]
    results = qdrant.query_points(collection_name=collection, query=query_vector, limit=top_k)
    return [{"file_path": hit.payload["file_path"], "score": round(hit.score, 4),
             "text_preview": hit.payload["text"][:120],
             "symbol_name": hit.payload.get("symbol_name", "n/a")}
            for hit in results.points]


def compare():
    """Compare naive retrieval against AST-aware retrieval on benchmark questions.

    Args:
        None.

    Returns:
        None. Results are printed for manual inspection.
    """
    questions = []
    with open(BENCHMARK_FILE) as f:
        for line in f:
            if line.strip():
                questions.append(json.loads(line))

    print(f"Comparing naive vs. AST-aware retrieval on {min(len(questions), 15)} questions\n")
    for q in questions[:15]:
        print(f"[{q['category']}] {q['question'][:70]}")
        naive = retrieve_from("anchor-repo-naive", q["question"])
        ast_aware = retrieve_from("anchor-repo-ast", q["question"])
        print(f"  Naive:     {', '.join(set(r['file_path'] for r in naive))}")
        print(f"  AST-aware: {', '.join(set(r['file_path'] for r in ast_aware))}")
        for r in ast_aware[:3]:
            print(f"    [{r['score']}] {r['symbol_name']} in {r['file_path']}")
        print()


if __name__ == "__main__":
    compare()