Галоўная / Артыкулы / Што такое RAG? Наўаты, гібрыдны, графічны і агентны патэрны

Што такое RAG? Наўаты, гібрыдны, графічны і агентны патэрны

Практычныя карточкі генеравання з адказамі на запыткі: класычны вектарны RAG, гібрыдны спосаб адказавання на запыткі, GraphRAG і агентны падходы — калі кожны з яных падходзіць і што ламаецца ў практычнай експлуатацыі.

5599 слоў

RAG (Retrieval-Augmented Generation) дае адказы за дапамогай зовнішняго источніка ведаць, а не толькі за дапамогай памяці модэлі. Спачатку выкарыстоўваюцца стосунаваныя фрагменты, якія пасля перадаюцца як контэкст, каб LLM могла практычна адпаведаць на запыт. Уявіце экзамен з адкрытым кнігам: модэль — это студент, а RAG сапраўджае кнігу. Адказы можна апдэйтаваць без падготовкі модэлі занова.

Як працюе RAG

Тры стадзіі:

  1. Індэксаванне — разбіўка дакументаў на часткі, якія можна шукаты (вектарны база дадзеных, граф, зворачаны індэкс і т. д.)
  2. Пошук — выявленне найболей падходзячых частак для запыту
  3. Генераванне — перадача эйх частак LLM для стварэння адказу

Якасць пошуку мае вялічынны ўплыв. Нават сильная модэль з слабым пошуком дае слабыя адказы; ключавая частка — спосаб выявлення контэксту.

1. Наўатывны RAG

Самы просты ўзор: уключаць дакументы, зберагчы вектары, адзначаць іх па сэмантычнай супаданосці пад час запиту.

Узор

Разбіваць дакументы на часткі, уключаць іх і зберагчы вектары. Пад час запиту уключаць сам запит, выбраць найбліжэйшыя суседні элементы і дазволіць модэлю напісаць адказ на адвяроте з тых частак.

Калі яго вжываць

  • Протатыпы і дэманстраціі
  • Маленькія, чыстыя корпусы (прыблізна менейш за ~50 тыс. частак)
  • Домены, дзе сэмантычная супаданосць пасвядчае намеру (запыты і адказы у прозе)

Калі ён перестае работаць

  • Точныя токены (імены, коды, ID-ы), якія не ўрахоўваюцца пад час уключэння
  • Няма фільтраў метадаў (дата, выканнік, катэгорыя)
  • Некалькасловы ранжыраванне без паўторнага ранжыравання
  • Полныя сканаванні, якія стаюць дорогімі пад час обробкі мільйонав частак

На шта звернуць увагу ў коде

Індексаванне ў ChromaDB, простыя запиты retrieve() на аднойчыні з косай схожасцю, а таксама системны запрос, який обмежвае модель наданым контекстам:

"""
Naive RAG — the simplest form of Retrieval-Augmented Generation.

Pattern:
  1. Chunk documents into small pieces
  2. Embed each chunk into a vector
  3. Store vectors in a vector database
  4. At query time: embed the query, find the closest chunk vectors
  5. Pass those chunks as context to the LLM

This is the baseline. It's quick to build and works well for small,
clean corpora where semantic similarity reliably maps to relevance.
"""

import os
from openai import OpenAI
import chromadb
from sentence_transformers import SentenceTransformer
from dotenv import load_dotenv

load_dotenv()


# ─────────────────────────────────────────────────────────────────
#  LLM abstraction — same code works with OpenAI, Groq, or Ollama
# ─────────────────────────────────────────────────────────────────

def get_llm():
    """
    Returns (client, model_name) for the configured provider.
    All three providers expose an OpenAI-compatible API, so the
    rest of the code doesn't need to change per provider.
    """
    provider = os.getenv("LLM_PROVIDER", "openai")

    if provider == "groq":
        client = OpenAI(
            api_key=os.getenv("GROQ_API_KEY"),
            base_url="https://api.groq.com/openai/v1"
        )
        return client, "llama-3.1-8b-instant"

    if provider == "ollama":
        # Ollama runs locally — no API key needed
        client = OpenAI(api_key="ollama", base_url="http://localhost:11434/v1")
        return client, os.getenv("OLLAMA_MODEL", "llama3.2")

    # Default: OpenAI
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    return client, "gpt-4o-mini"


# ─────────────────────────────────────────────────────────────────
#  Knowledge base — the documents we want to search over
# ─────────────────────────────────────────────────────────────────

# In a real system, these would come from PDFs, databases, or APIs.
# Here we use a small set of space exploration facts so you can run
# this without any external data source.
DOCUMENTS = [
    "The Apollo 11 mission landed the first humans on the Moon on July 20, 1969. "
    "Neil Armstrong and Buzz Aldrin walked on the surface while Michael Collins orbited above.",

    "The James Webb Space Telescope (JWST) launched on December 25, 2021. "
    "It observes in infrared light, letting it see through dust clouds and capture images of the earliest galaxies.",

    "Mars has two small moons called Phobos and Deimos. "
    "Scientists believe they are captured asteroids, not moons that formed alongside the planet.",

    "Voyager 1 is the most distant human-made object ever launched. "
    "It entered interstellar space in 2012 and continues to transmit data back to Earth from over 23 billion km away.",

    "SpaceX's Falcon 9 is a partially reusable rocket. "
    "The first stage booster autonomously lands back on Earth or a drone ship and is refurbished for future flights.",

    "The International Space Station (ISS) orbits Earth at ~400 km altitude and travels at roughly 28,000 km/h. "
    "It has been continuously inhabited since November 2000.",

    "Saturn's rings are made mostly of water ice and rocky debris, ranging from microscopic grains to chunks the size of a house. "
    "Despite spanning hundreds of thousands of kilometers, the rings are only about 10 meters thick in some places.",

    "The Hubble Space Telescope has been operating since 1990. "
    "Its deep-field images revealed thousands of galaxies in a patch of sky that appeared completely empty to the naked eye.",
]


# ─────────────────────────────────────────────────────────────────
#  Step 1 — Embed and index all documents
# ─────────────────────────────────────────────────────────────────

# sentence-transformers runs locally — no embedding API key required.
# all-MiniLM-L6-v2 is small (~80MB) but good enough for most tasks.
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

# ChromaDB in-memory client — no disk, no server, no setup.
# For production, swap this for chromadb.PersistentClient() or a hosted DB.
chroma_client = chromadb.Client()
collection = chroma_client.create_collection("space_facts")

print("📥 Indexing documents...")
doc_embeddings = embedding_model.encode(DOCUMENTS).tolist()

collection.add(
    documents=DOCUMENTS,
    embeddings=doc_embeddings,
    ids=[f"doc_{i}" for i in range(len(DOCUMENTS))]
)
print(f"✅ Indexed {len(DOCUMENTS)} documents.\n")


# ─────────────────────────────────────────────────────────────────
#  Step 2 — Retrieve relevant chunks for a query
# ─────────────────────────────────────────────────────────────────

def retrieve(query: str, top_k: int = 3) -> list[str]:
    """
    Embeds the query and finds the top-k most semantically similar chunks.

    This is pure cosine similarity in vector space — there's no keyword
    matching, no filters, no reranking. Simple, but fragile at scale.
    """
    query_embedding = embedding_model.encode([query]).tolist()

    results = collection.query(
        query_embeddings=query_embedding,
        n_results=top_k
    )

    # results["documents"] is a list-of-lists (one per query),
    # so we take index [0] for our single query
    return results["documents"][0]


# ─────────────────────────────────────────────────────────────────
#  Step 3 — Generate an answer using the retrieved context
# ─────────────────────────────────────────────────────────────────

def generate(query: str, context_chunks: list[str]) -> str:
    """
    Passes the query + retrieved chunks to the LLM.

    The system prompt instructs the LLM to answer ONLY from the provided
    context. Without this constraint the model might blend retrieved facts
    with its own training data, which defeats the purpose of RAG.
    """
    client, model = get_llm()

    # Join chunks into a numbered list for readability
    context = "\n".join(f"{i+1}. {chunk}" for i, chunk in enumerate(context_chunks))

    messages = [
        {
            "role": "system",
            "content": (
                "You are a factual assistant. Answer the user's question using ONLY "
                "the context provided. If the answer is not in the context, say "
                "'I don't have that information.'"
            )
        },
        {
            "role": "user",
            "content": f"Context:\n{context}\n\nQuestion: {query}"
        }
    ]

    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0  # deterministic — we want factual answers, not creative ones
    )

    return response.choices[0].message.content


# ─────────────────────────────────────────────────────────────────
#  Run the full RAG pipeline on a few example queries
# ─────────────────────────────────────────────────────────────────

queries = [
    "When did the first humans land on the Moon?",
    "How fast does the ISS travel?",
    "What are Saturn's rings made of?",
    "Who invented the telephone?",  # not in our knowledge base — watch what happens
]

for query in queries:
    print(f"❓ {query}")

    chunks = retrieve(query)
    print(f"   Retrieved chunks:")
    for chunk in chunks:
        # Print just the first 90 chars so the output stays readable
        print(f"     • {chunk[:90]}...")

    answer = generate(query, chunks)
    print(f"   💬 {answer}\n")

2. Гібрыдны RAG

У практычных запитах зазвычай выкарыстоўваюцца вектарны пошук, ключавыя словы BM25, фільтры метаданых і методы переранкавання з викорыстаннем крос-кодэра.

Проблемы простага RAG

Чыстыя вектары добра справляюцца з запитамі, спрямованымі на розумэнне значэння, але маюць труднасці з точным падборам, фільтраванымі даннымі, такімі як “толькі юрыдычныя дакументы 2023 года”, і “шумнымі” суседнімі элементамі. Ключавыя словы, метаданы і методы переранкавання дапамагаюць заполніць гэтыя праслейкі.

Патерн

Query
  │
  ├─→ Metadata filter  ──→ narrow candidate pool
  │
  ├─→ Vector search    ──┐
  │                      ├─→ RRF fusion  ──→ Cross-encoder rerank  ──→ Top-k  ──→ LLM
  └─→ BM25 search     ──┘

Этап 1 — Фільтраўанне метаданых звужае колькасць кандыдатаў за дапамогою структураваных поль пры перадачы данных на эмбеддінг.

Этап 2 — Гібрыдны пошук адбываецца паралельна: выкарыстоўваюцца як вектарная схожасць, так і метод BM25, ў результате чаго як семантыка, так і точныя тэрміны павышаюць свой ранг.

Этап 3 — сяроджэнне RRF спаўнае рангаваныя спискі за дапамою методу Reciprocal Rank Fusion (score = 1 / (k + rank)) без нормалізацыі несумесных значэнняў рангу.

Этап 4 — паўторнае рангаванне за дапамою крос-кодэра адзіночна оцінюе пары (query, document) для атрымання точнага фінальнага порядку — гэта медленней, але даходзіць да болей выражаных рэзультатаў.

Калі яго вжываць

  • Праблемы выкарыстоўвання ў практыцы, калі неабходныя як воспамінанне, так і тачнасць
  • Корпусы з корисным метаданным
  • Запиты, якія ўключаюць як семантычныя, так і ключавы элементы
  • Случаі, калі просты метод RAG не можа знайсці неабходную інфармацыю або паспелаў заняць весь контэкстны прызмат
"""
Hybrid RAG — vector search + BM25 keyword search + metadata filtering + reranking.

The Problem with Naive RAG:
  Pure vector similarity misses exact keyword matches (names, codes, product IDs)
  and has no way to filter by structured metadata (date, source, category).

The Fix — three upgrades layered on top of each other:
  1. Metadata filtering  → narrow the search space before any ranking
  2. Hybrid search       → run vector search AND BM25 in parallel, fuse scores
  3. Reranking           → use a cross-encoder to re-score the fused candidates

This combination is what most production RAG systems use.
"""

import os
import math
from openai import OpenAI
from rank_bm25 import BM25Okapi
import chromadb
from sentence_transformers import SentenceTransformer, CrossEncoder
from dotenv import load_dotenv

load_dotenv()


# ─────────────────────────────────────────────────────────────────
#  LLM abstraction — same pattern as 01-naive-rag
# ─────────────────────────────────────────────────────────────────

def get_llm():
    provider = os.getenv("LLM_PROVIDER", "openai")
    if provider == "groq":
        return OpenAI(
            api_key=os.getenv("GROQ_API_KEY"),
            base_url="https://api.groq.com/openai/v1"
        ), "llama-3.1-8b-instant"
    if provider == "ollama":
        return OpenAI(api_key="ollama", base_url="http://localhost:11434/v1"), \
               os.getenv("OLLAMA_MODEL", "llama3.2")
    return OpenAI(api_key=os.getenv("OPENAI_API_KEY")), "gpt-4o-mini"


# ─────────────────────────────────────────────────────────────────
#  Knowledge base — documents WITH metadata
# ─────────────────────────────────────────────────────────────────

# Each document now carries structured metadata alongside the text.
# This is the foundation for metadata filtering — you can pre-filter
# before running any expensive embedding or BM25 computation.
DOCUMENTS = [
    {
        "id": "doc_0",
        "text": "The Apollo 11 mission landed the first humans on the Moon on July 20, 1969. "
                "Neil Armstrong and Buzz Aldrin walked on the lunar surface.",
        "metadata": {"category": "missions", "year": 1969, "source": "nasa.gov"}
    },
    {
        "id": "doc_1",
        "text": "The James Webb Space Telescope (JWST) launched on December 25, 2021. "
                "It uses infrared imaging to observe early universe galaxies and exoplanet atmospheres.",
        "metadata": {"category": "telescopes", "year": 2021, "source": "nasa.gov"}
    },
    {
        "id": "doc_2",
        "text": "SpaceX's Falcon 9 is a partially reusable rocket. The first stage booster lands "
                "autonomously after launch and is refurbished for reuse.",
        "metadata": {"category": "rockets", "year": 2015, "source": "spacex.com"}
    },
    {
        "id": "doc_3",
        "text": "NASA's Artemis program aims to return humans to the Moon by 2026 and establish "
                "a sustainable lunar presence as a stepping stone to Mars.",
        "metadata": {"category": "missions", "year": 2022, "source": "nasa.gov"}
    },
    {
        "id": "doc_4",
        "text": "The Hubble Space Telescope has been operating since 1990, producing iconic images "
                "of nebulae and distant galaxies in visible and ultraviolet light.",
        "metadata": {"category": "telescopes", "year": 1990, "source": "nasa.gov"}
    },
    {
        "id": "doc_5",
        "text": "SpaceX's Starship is designed for full reusability — both the Super Heavy booster "
                "and the Starship upper stage return and land after flight.",
        "metadata": {"category": "rockets", "year": 2023, "source": "spacex.com"}
    },
    {
        "id": "doc_6",
        "text": "The Mars Perseverance Rover landed in Jezero Crater in February 2021. "
                "It is searching for signs of ancient microbial life and collecting rock samples.",
        "metadata": {"category": "missions", "year": 2021, "source": "nasa.gov"}
    },
    {
        "id": "doc_7",
        "text": "The Nancy Grace Roman Space Telescope is NASA's next flagship telescope, "
                "designed to survey wide fields of the sky for dark energy and exoplanets.",
        "metadata": {"category": "telescopes", "year": 2026, "source": "nasa.gov"}
    },
]


# ─────────────────────────────────────────────────────────────────
#  Build the two search indexes
# ─────────────────────────────────────────────────────────────────

# --- Embedding model (vector search) ---
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

# --- Cross-encoder (reranking) ---
# A cross-encoder takes (query, document) pairs and outputs a relevance score.
# It's slower than bi-encoders but far more accurate — used as a final step.
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

# --- ChromaDB (vector store) ---
chroma_client = chromadb.Client()
collection = chroma_client.create_collection("space_hybrid")

texts = [d["text"] for d in DOCUMENTS]
doc_embeddings = embedding_model.encode(texts).tolist()

collection.add(
    documents=texts,
    embeddings=doc_embeddings,
    ids=[d["id"] for d in DOCUMENTS],
    metadatas=[d["metadata"] for d in DOCUMENTS]
)

# --- BM25 (keyword search) ---
# BM25 (Best Match 25) is a classic probabilistic ranking function.
# It scores documents based on term frequency and inverse document frequency.
# It excels at exact keyword matches that embeddings often miss.
tokenized_corpus = [doc.lower().split() for doc in texts]
bm25_index = BM25Okapi(tokenized_corpus)

print(f"✅ Indexed {len(DOCUMENTS)} documents (vector + BM25).\n")


# ─────────────────────────────────────────────────────────────────
#  Retrieval — 3-stage pipeline
# ─────────────────────────────────────────────────────────────────

def filter_by_metadata(category: str | None = None, min_year: int | None = None) -> list[int]:
    """
    Stage 1: Metadata filtering.

    Returns the indices of documents that match the given filters.
    Filtering BEFORE vector/keyword search narrows the corpus, making
    retrieval faster and more precise — especially important at scale.
    """
    indices = []
    for i, doc in enumerate(DOCUMENTS):
        meta = doc["metadata"]
        if category and meta["category"] != category:
            continue
        if min_year and meta["year"] < min_year:
            continue
        indices.append(i)
    return indices


def reciprocal_rank_fusion(
    vector_ids: list[str],
    bm25_indices: list[int],
    k: int = 60
) -> list[str]:
    """
    Stage 2b: Score fusion using Reciprocal Rank Fusion (RRF).

    RRF combines rankings from multiple retrieval systems without needing
    to normalize their raw scores. The formula is: score = 1 / (k + rank).
    Results ranked high in EITHER system get boosted — catching what the
    other system missed.
    """
    scores: dict[str, float] = {}

    # Score from vector ranking
    for rank, doc_id in enumerate(vector_ids):
        scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)

    # Score from BM25 ranking
    for rank, doc_idx in enumerate(bm25_indices):
        doc_id = DOCUMENTS[doc_idx]["id"]
        scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)

    # Sort by combined score, highest first
    return sorted(scores, key=scores.get, reverse=True)


def hybrid_retrieve(
    query: str,
    category: str | None = None,
    min_year: int | None = None,
    top_k: int = 3,
    candidate_k: int = 6
) -> list[dict]:
    """
    Full hybrid retrieval:
      1. Filter by metadata  → narrow the candidate pool
      2. Vector search       → semantic similarity
      3. BM25 search         → keyword matching
      4. RRF fusion          → combine the two rankings
      5. Cross-encoder rerank → final high-precision scoring

    candidate_k: how many candidates to gather before reranking.
                 More candidates = more recall, but slower reranking.
    """

    # ── Stage 1: Metadata filter ──────────────────────────────────
    allowed_indices = filter_by_metadata(category=category, min_year=min_year)
    allowed_ids = {DOCUMENTS[i]["id"] for i in allowed_indices}
    allowed_texts = [DOCUMENTS[i]["text"] for i in allowed_indices]

    if not allowed_ids:
        print("  ⚠️  No documents match the metadata filters.")
        return []

    # ── Stage 2a: Vector search (within filtered set) ─────────────
    query_embedding = embedding_model.encode([query]).tolist()
    vector_results = collection.query(
        query_embeddings=query_embedding,
        n_results=min(candidate_k, len(allowed_ids)),
        where={"category": category} if category else None  # ChromaDB metadata filter
    )
    vector_ids = vector_results["documents"][0] and vector_results["ids"][0] or []
    # Keep only IDs that passed our metadata filter
    vector_ids = [id_ for id_ in vector_ids if id_ in allowed_ids]

    # ── Stage 2b: BM25 keyword search (within filtered set) ───────
    tokenized_query = query.lower().split()
    # Score all filtered documents with BM25
    bm25_filtered = BM25Okapi([DOCUMENTS[i]["text"].lower().split() for i in allowed_indices])
    bm25_scores = bm25_filtered.get_scores(tokenized_query)
    # Get indices sorted by score (within the filtered subset)
    top_bm25_local = sorted(range(len(allowed_indices)), key=lambda i: bm25_scores[i], reverse=True)[:candidate_k]
    # Map back to original DOCUMENTS indices
    top_bm25_global = [allowed_indices[i] for i in top_bm25_local]

    # ── Stage 3: RRF fusion ────────────────────────────────────────
    fused_ids = reciprocal_rank_fusion(vector_ids, top_bm25_global)[:candidate_k]

    # ── Stage 4: Cross-encoder reranking ──────────────────────────
    # Fetch the actual text for each candidate
    candidates = [
        next(d for d in DOCUMENTS if d["id"] == doc_id)
        for doc_id in fused_ids
        if any(d["id"] == doc_id for d in DOCUMENTS)
    ]

    # Cross-encoder scores each (query, document) pair independently
    # This is more accurate than embedding similarity but O(n) slower
    pairs = [(query, c["text"]) for c in candidates]
    rerank_scores = reranker.predict(pairs)

    # Sort candidates by reranker score, take top_k
    ranked = sorted(zip(candidates, rerank_scores), key=lambda x: x[1], reverse=True)
    return [doc for doc, _ in ranked[:top_k]]


# ─────────────────────────────────────────────────────────────────
#  Generate
# ─────────────────────────────────────────────────────────────────

def generate(query: str, context_docs: list[dict]) -> str:
    client, model = get_llm()

    context = "\n".join(
        f"{i+1}. [{doc['metadata']['category']} | {doc['metadata']['year']}] {doc['text']}"
        for i, doc in enumerate(context_docs)
    )

    messages = [
        {
            "role": "system",
            "content": "You are a factual assistant. Answer using ONLY the provided context."
        },
        {
            "role": "user",
            "content": f"Context:\n{context}\n\nQuestion: {query}"
        }
    ]

    response = client.chat.completions.create(
        model=model, messages=messages, temperature=0
    )
    return response.choices[0].message.content


# ─────────────────────────────────────────────────────────────────
#  Run example queries
# ─────────────────────────────────────────────────────────────────

examples = [
    {
        "query": "Tell me about reusable rockets",
        "filters": {"category": "rockets"},          # only look at rocket documents
        "description": "Metadata filter: rockets only"
    },
    {
        "query": "What telescope launched most recently?",
        "filters": {"category": "telescopes", "min_year": 2020},  # telescopes after 2020
        "description": "Metadata filter: telescopes after 2020"
    },
    {
        "query": "NASA Mars mission 2021",
        "filters": {"category": "missions"},
        "description": "BM25 shines: exact keywords (NASA, Mars, 2021)"
    },
]

for ex in examples:
    print(f"❓ {ex['query']}")
    print(f"   Filter: {ex['description']}")

    docs = hybrid_retrieve(
        query=ex["query"],
        category=ex["filters"].get("category"),
        min_year=ex["filters"].get("min_year"),
    )

    print(f"   Retrieved {len(docs)} docs after reranking:")
    for doc in docs:
        print(f"     • [{doc['metadata']['category']} | {doc['metadata']['year']}] {doc['text'][:80]}...")

    answer = generate(ex["query"], docs)
    print(f"   💬 {answer}\n")

3. Graph RAG

У гэтым падходзе граф знанняяў выступае як основа для пошуку (альбо разам, альбо пасля вектараў).

Проблема, якую рашае GraphRAG

Падобнасць спрашва: «Какі тэксты належаць да майго запиту?» Запытанні пра звязкі — «Хто інвеставаў у OpenAI?», «Што купіла Microsoft?», «Як з’яўляецца звязак між X і Y?» — ёсць пераходы па графе. Эмбеддінгі спрощуюць структуру; графы залучаюць типаваныя канцэнтры як частку сваей структуры.

Шаблон

Documents / structured data
  │
  └─→ Entity + Relationship extraction  ──→  Knowledge Graph (nodes + edges)

Query
  │
  ├─→ Entity extraction  ──→  Find matching nodes
  │
  └─→ Graph traversal (N hops)  ──→  Subgraph context  ──→  LLM  ──→  Answer
"""
GraphRAG — retrieval using a knowledge graph instead of (or alongside) a vector DB.

Why a graph?
  Documents store text. Graphs store *relationships* — who owns what, what caused what,
  which thing is part of which system. When your questions are about connections
  ("Who invested in OpenAI?", "What products did Microsoft acquire?"), traversing a
  graph gives you exactly the right context without keyword or semantic guesswork.

Pattern:
  1. Build a knowledge graph from entities and their relationships
  2. Extract entities from the query (simple keyword matching here; NER in production)
  3. Find matching nodes in the graph
  4. Traverse 1-2 hops to collect related context
  5. Format the subgraph as text and pass to the LLM
"""

import os
import networkx as nx
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()


# ─────────────────────────────────────────────────────────────────
#  LLM abstraction
# ─────────────────────────────────────────────────────────────────

def get_llm():
    provider = os.getenv("LLM_PROVIDER", "openai")
    if provider == "groq":
        return OpenAI(
            api_key=os.getenv("GROQ_API_KEY"),
            base_url="https://api.groq.com/openai/v1"
        ), "llama-3.1-8b-instant"
    if provider == "ollama":
        return OpenAI(api_key="ollama", base_url="http://localhost:11434/v1"), \
               os.getenv("OLLAMA_MODEL", "llama3.2")
    return OpenAI(api_key=os.getenv("OPENAI_API_KEY")), "gpt-4o-mini"


# ─────────────────────────────────────────────────────────────────
#  Knowledge Graph — entities and relationships in the AI industry
# ─────────────────────────────────────────────────────────────────

# We use networkx's directed graph: nodes are entities, edges are relationships.
# In a production GraphRAG system, this graph would be extracted automatically
# from documents using an LLM or NER pipeline. Here we build it manually.
G = nx.DiGraph()

# ── Nodes — each entity has a type and a short description ───────
entities = [
    ("OpenAI",          {"type": "company",  "desc": "AI research company, creator of GPT-4, DALL-E, and Sora"}),
    ("Microsoft",       {"type": "company",  "desc": "Technology giant, owner of Azure, Office, and GitHub"}),
    ("Google",          {"type": "company",  "desc": "Technology company, owner of Search, YouTube, and DeepMind"}),
    ("DeepMind",        {"type": "company",  "desc": "AI research lab, creator of AlphaFold and Gemini"}),
    ("Anthropic",       {"type": "company",  "desc": "AI safety company, creator of the Claude model family"}),
    ("Sam Altman",      {"type": "person",   "desc": "CEO of OpenAI, former president of Y Combinator"}),
    ("Ilya Sutskever",  {"type": "person",   "desc": "Co-founder of OpenAI and SSI, key architect of GPT models"}),
    ("Demis Hassabis",  {"type": "person",   "desc": "CEO and co-founder of DeepMind"}),
    ("Dario Amodei",    {"type": "person",   "desc": "CEO and co-founder of Anthropic, former VP Research at OpenAI"}),
    ("GPT-4",           {"type": "product",  "desc": "Large language model by OpenAI, released in 2023"}),
    ("Claude",          {"type": "product",  "desc": "LLM family by Anthropic, focused on safety and helpfulness"}),
    ("Gemini",          {"type": "product",  "desc": "Multimodal LLM by Google DeepMind, released in 2023"}),
    ("AlphaFold",       {"type": "product",  "desc": "AI system by DeepMind that predicts protein structures"}),
    ("Azure",           {"type": "product",  "desc": "Microsoft's cloud platform, hosts OpenAI models via Azure OpenAI"}),
    ("GitHub",          {"type": "product",  "desc": "Code hosting platform owned by Microsoft"}),
    ("Y Combinator",    {"type": "org",      "desc": "Startup accelerator that funded companies like OpenAI, Dropbox, Airbnb"}),
]

G.add_nodes_from(entities)

# ── Edges — directional relationships between entities ────────────
# Format: (source, target, {"relation": "...", "detail": "..."})
relationships = [
    ("Sam Altman",     "OpenAI",       {"relation": "CEO_of",       "detail": "CEO since 2019 (with a brief ouster in 2023)"}),
    ("Ilya Sutskever", "OpenAI",       {"relation": "co_founded",   "detail": "Co-founded OpenAI in 2015 alongside Sam Altman"}),
    ("Dario Amodei",   "Anthropic",    {"relation": "co_founded",   "detail": "Founded Anthropic in 2021 after leaving OpenAI"}),
    ("Dario Amodei",   "OpenAI",       {"relation": "former_VP_at", "detail": "Was VP of Research at OpenAI before leaving"}),
    ("Demis Hassabis", "DeepMind",     {"relation": "co_founded",   "detail": "Co-founded DeepMind in 2010, acquired by Google in 2014"}),
    ("Microsoft",      "OpenAI",       {"relation": "invested_in",  "detail": "Invested ~$13B in OpenAI across multiple rounds"}),
    ("Google",         "Anthropic",    {"relation": "invested_in",  "detail": "Invested ~$300M in Anthropic in 2023"}),
    ("Google",         "DeepMind",     {"relation": "acquired",     "detail": "Acquired DeepMind in 2014 for ~$500M"}),
    ("Microsoft",      "GitHub",       {"relation": "acquired",     "detail": "Acquired GitHub in 2018 for $7.5B"}),
    ("OpenAI",         "GPT-4",        {"relation": "created",      "detail": "Released GPT-4 in March 2023"}),
    ("Anthropic",      "Claude",       {"relation": "created",      "detail": "Claude 3 family released in 2024"}),
    ("DeepMind",       "Gemini",       {"relation": "created",      "detail": "Gemini 1.0 released December 2023 as GPT-4 competitor"}),
    ("DeepMind",       "AlphaFold",    {"relation": "created",      "detail": "AlphaFold 2 solved protein structure prediction in 2020"}),
    ("Microsoft",      "Azure",        {"relation": "owns",         "detail": "Azure hosts OpenAI's models via Azure OpenAI Service"}),
    ("OpenAI",         "Azure",        {"relation": "partners_with","detail": "OpenAI's API and models are available through Azure"}),
    ("Sam Altman",     "Y Combinator", {"relation": "led",          "detail": "Was president of Y Combinator from 2014 to 2019"}),
]

G.add_edges_from(relationships)

print(f"✅ Knowledge graph built: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges.\n")


# ─────────────────────────────────────────────────────────────────
#  Step 1 — Entity extraction from the query
# ─────────────────────────────────────────────────────────────────

def extract_entities(query: str) -> list[str]:
    """
    Find which graph nodes are mentioned in the query.

    This is simple case-insensitive substring matching — good enough for demos.
    In production you'd use an LLM or a NER model to extract entities, which
    handles synonyms, abbreviations, and entities not spelled out exactly.
    """
    query_lower = query.lower()
    return [node for node in G.nodes if node.lower() in query_lower]


# ─────────────────────────────────────────────────────────────────
#  Step 2 — Graph traversal to collect context
# ─────────────────────────────────────────────────────────────────

def retrieve_from_graph(query: str, max_hops: int = 2) -> str:
    """
    Retrieves context by traversing the graph from query entities.

    For each entity found in the query:
      - Include the entity's own description
      - Walk outgoing edges (what this entity does / relates to)
      - Walk incoming edges (what points to this entity)
      - Optionally go 1 more hop for richer context

    max_hops=1 gives direct neighbors; max_hops=2 also includes neighbors' neighbors.
    More hops = richer context but also more noise.
    """
    seed_entities = extract_entities(query)

    if not seed_entities:
        return "No matching entities found in the knowledge graph."

    print(f"   Entities found: {seed_entities}")

    context_lines = []
    visited = set()

    def collect_node(node: str):
        """Collect a node's description and all its edges as text."""
        if node in visited:
            return
        visited.add(node)

        data = G.nodes[node]
        context_lines.append(f"[{data['type'].upper()}] {node}: {data['desc']}")

        # Outgoing edges: what this entity does/relates to
        for _, target, edge_data in G.out_edges(node, data=True):
            context_lines.append(
                f"  → {node} --[{edge_data['relation']}]--> {target}: {edge_data['detail']}"
            )

        # Incoming edges: what points to this entity
        for source, _, edge_data in G.in_edges(node, data=True):
            context_lines.append(
                f"  ← {source} --[{edge_data['relation']}]--> {node}: {edge_data['detail']}"
            )

    # Collect seed entities and their neighbors up to max_hops
    for entity in seed_entities:
        collect_node(entity)

        if max_hops >= 2:
            # Also collect 1-hop neighbors for richer context
            for neighbor in list(G.successors(entity)) + list(G.predecessors(entity)):
                collect_node(neighbor)

    return "\n".join(context_lines)


# ─────────────────────────────────────────────────────────────────
#  Step 3 — Generate answer from graph context
# ─────────────────────────────────────────────────────────────────

def generate(query: str, graph_context: str) -> str:
    client, model = get_llm()

    messages = [
        {
            "role": "system",
            "content": (
                "You are a helpful assistant with access to a knowledge graph about the AI industry. "
                "Answer the question using ONLY the graph context provided. "
                "Be concise and factual."
            )
        },
        {
            "role": "user",
            "content": f"Knowledge Graph Context:\n{graph_context}\n\nQuestion: {query}"
        }
    ]

    response = client.chat.completions.create(
        model=model, messages=messages, temperature=0
    )
    return response.choices[0].message.content


# ─────────────────────────────────────────────────────────────────
#  Run example queries
# ─────────────────────────────────────────────────────────────────

queries = [
    "Who co-founded OpenAI?",
    "What has Microsoft invested in or acquired?",
    "What is the relationship between Google and DeepMind?",
    "What did Dario Amodei do before Anthropic?",
]

for query in queries:
    print(f"❓ {query}")

    context = retrieve_from_graph(query)
    print(f"   Graph context snippet:\n     {context[:200]}...")

    answer = generate(query, context)
    print(f"   💬 {answer}\n")

Структура графа

  • Вузлы — аб’екты (компаніі, людзі, продукты) з апісамі
  • Канцэнтры — типаваныя, направленыя звязкі (invested_in, acquired, CEO_of, …)

Адзысканне через пераходы

  1. Выявіце аб’екты, упомянутыя ў запытанні
  2. Выкорыстаце атрыбуты тых вузлаў і ўпрывязаныя канцэнтры як выхадную пункт
  3. Па выбору расшырьце пераходы колькіма шагамі, ўбачлівейшы такім чынам контекст

Калі яго вжываць

  • Калі є чыстае апісанне аб’ектаў і названных звязкаў
  • Пытанні на кшталт «Хто/чы ўзьязначаны з…?» / «Што X зрабіў з Y?»
  • Паходы, якія можна паказаць корыстнікаў, і якія ўсё можна поясніць
  • Знанні, якіе апдэйтуюцца без павольнага перыявлення ўсіх дадзеных
  • 4. Agentic RAG

    Модель плануе процес выкарыстоўвання дадзеных: якія інструменты, у какой спраўе, і калі ёсць достатнек доказаў.

    Проблема фіксаваных пайплайнаў

    Наўкія, гібрыдныя та графавыя пайплайны закодуюць адзін паход для кожнага пытання. Неодназначныя або багатоэтапныя запиты выклікаюць патрэбу ў розгалужэннях.

    Шаблон — цыкл ReAct

    ReAct = Разумовыя вычынкі + Дзеяння: агент думае, вызывае інструмент для выкарыстоўвання дадзеных аб вычыслаў, спазірае на рэзультаты і павторюе процес, пакуль не зможа даўчы.

    """
    Agentic RAG — the LLM plans its own retrieval strategy.
    
    The Problem with Fixed Pipelines:
      In Naive/Hybrid/GraphRAG, the retrieval strategy is hardcoded. Every query
      goes through the same pipeline regardless of what the question actually needs.
    
      But questions vary:
        "What is SpaceX?" → single lookup, done
        "Compare SpaceX and NASA's approach to Mars" → needs two retrievals + synthesis
        "What is the orbital speed of the ISS in miles per hour?" → needs retrieval + math
    
      A static pipeline can't adapt. An agent can.
    
    The Agentic Approach:
      We implement a ReAct loop (Reason + Act) where the LLM:
        1. Reads the question
        2. Decides which tool to use and with what input
        3. Sees the tool's output
        4. Decides whether it has enough to answer, or needs another tool call
        5. Repeats until it can give a final answer
    
    Tools available to the agent:
      - search_docs(query)     → vector search over a document corpus
      - get_entity(name)       → exact lookup of a known entity (like a mini-graph)
      - calculate(expression)  → evaluate a math expression safely
    """
    
    import os
    import json
    import ast
    import operator
    import chromadb
    from openai import OpenAI
    from sentence_transformers import SentenceTransformer
    from dotenv import load_dotenv
    
    load_dotenv()
    
    
    # ─────────────────────────────────────────────────────────────────
    #  LLM abstraction
    # ─────────────────────────────────────────────────────────────────
    
    def get_llm():
        provider = os.getenv("LLM_PROVIDER", "openai")
        if provider == "groq":
            return OpenAI(
                api_key=os.getenv("GROQ_API_KEY"),
                base_url="https://api.groq.com/openai/v1"
            ), "llama-3.1-8b-instant"
        if provider == "ollama":
            return OpenAI(api_key="ollama", base_url="http://localhost:11434/v1"), \
                   os.getenv("OLLAMA_MODEL", "llama3.2")
        return OpenAI(api_key=os.getenv("OPENAI_API_KEY")), "gpt-4o-mini"
    
    
    # ─────────────────────────────────────────────────────────────────
    #  Knowledge base — indexed for search_docs tool
    # ─────────────────────────────────────────────────────────────────
    
    DOCUMENTS = [
        "SpaceX was founded in 2002 by Elon Musk with the goal of making space travel cheaper "
        "and eventually colonizing Mars. Its Falcon 9 is the world's first orbital-class reusable rocket.",
    
        "NASA was founded in 1958 and is a US government agency. Its Mars missions include "
        "Curiosity (2012) and Perseverance (2021). NASA's SLS rocket is expendable, not reusable.",
    
        "The International Space Station (ISS) orbits Earth at 408 km altitude and travels "
        "at 7.66 km/s (27,576 km/h or 17,132 mph). It completes one orbit every 92 minutes.",
    
        "SpaceX's Starship is a fully reusable spacecraft designed for Mars missions, lunar landings, "
        "and point-to-point Earth travel. Its first successful orbital flight was in 2024.",
    
        "The Artemis program is NASA's plan to return humans to the Moon. Artemis 1 (2022) was "
        "uncrewed. Artemis 2 (2024) will be the first crewed test. Moon landing planned for Artemis 3.",
    
        "Mars is approximately 225 million km from Earth on average. A one-way trip with current "
        "propulsion takes 7-9 months. SpaceX aims to cut this with Starship's higher thrust-to-weight ratio.",
    
        "Blue Origin is a space company founded by Jeff Bezos in 2000. Its New Shepard rocket does "
        "suborbital tourism flights. Its New Glenn rocket is designed for orbital launches.",
    
        "Orbital velocity is the minimum speed needed to stay in orbit. At 400 km altitude, "
        "this is approximately 7.66 km/s. Below this, the spacecraft would fall back to Earth.",
    ]
    
    # Build in-memory vector index for the search_docs tool
    embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
    chroma_client = chromadb.Client()
    collection = chroma_client.create_collection("agent_kb")
    collection.add(
        documents=DOCUMENTS,
        embeddings=embedding_model.encode(DOCUMENTS).tolist(),
        ids=[f"doc_{i}" for i in range(len(DOCUMENTS))]
    )
    
    # Entity lookup table for the get_entity tool
    ENTITIES = {
        "spacex":    "SpaceX: Private space company by Elon Musk. Falcon 9 (reusable), Starship (Mars). Founded 2002.",
        "nasa":      "NASA: US government space agency. Artemis (Moon), Perseverance (Mars). Founded 1958.",
        "iss":       "ISS: Orbits at 408km, speed 7.66 km/s (17,132 mph), 92-min orbit. Inhabited since 2000.",
        "blue origin": "Blue Origin: Jeff Bezos' space company. New Shepard (suborbital), New Glenn (orbital). Founded 2000.",
        "starship":  "Starship: SpaceX's fully reusable Mars rocket. First orbital flight 2024. Largest rocket ever built.",
    }
    
    print("✅ Agent knowledge base ready.\n")
    
    
    # ─────────────────────────────────────────────────────────────────
    #  Tool definitions
    # ─────────────────────────────────────────────────────────────────
    
    # Tool specs in OpenAI function-calling format.
    # The agent (LLM) reads these to know what tools exist and how to call them.
    TOOLS = [
        {
            "type": "function",
            "function": {
                "name": "search_docs",
                "description": "Search the knowledge base using semantic similarity. Use for general questions about space, rockets, missions, or orbit.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "The search query"}
                    },
                    "required": ["query"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_entity",
                "description": "Look up a specific named entity (SpaceX, NASA, ISS, Blue Origin, Starship). Use when you need precise facts about a known entity.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string", "description": "Entity name to look up"}
                    },
                    "required": ["name"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "calculate",
                "description": "Evaluate a mathematical expression. Use for unit conversions or arithmetic. Example: '7.66 * 3600' to convert km/s to km/h.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "expression": {"type": "string", "description": "A safe arithmetic expression using +, -, *, /, (, ), and numbers"}
                    },
                    "required": ["expression"]
                }
            }
        }
    ]
    
    
    # ─────────────────────────────────────────────────────────────────
    #  Tool implementations
    # ─────────────────────────────────────────────────────────────────
    
    def search_docs(query: str) -> str:
        """Semantic search over the knowledge base."""
        query_embedding = embedding_model.encode([query]).tolist()
        results = collection.query(query_embeddings=query_embedding, n_results=2)
        chunks = results["documents"][0]
        return "\n---\n".join(chunks)
    
    
    def get_entity(name: str) -> str:
        """Exact entity lookup."""
        result = ENTITIES.get(name.lower().strip())
        return result if result else f"No entity found for '{name}'. Try: {list(ENTITIES.keys())}"
    
    
    def calculate(expression: str) -> str:
        """
        Safe arithmetic evaluator — no eval(), no exec().
        Only allows numbers and basic operators to prevent code injection.
        """
        # Only allow digits, spaces, and basic math operators
        allowed_chars = set("0123456789 +-*/().")
        if not all(c in allowed_chars for c in expression):
            return "Error: Only basic arithmetic is allowed (+, -, *, /, parentheses, numbers)."
        try:
            # Parse to AST and evaluate node-by-node — never calls eval()
            tree = ast.parse(expression, mode="eval")
            ops = {
                ast.Add: operator.add, ast.Sub: operator.sub,
                ast.Mult: operator.mul, ast.Div: operator.truediv,
                ast.USub: operator.neg
            }
    
            def eval_node(node):
                if isinstance(node, ast.Constant):
                    return node.value
                if isinstance(node, ast.BinOp):
                    return ops[type(node.op)](eval_node(node.left), eval_node(node.right))
                if isinstance(node, ast.UnaryOp):
                    return ops[type(node.op)](eval_node(node.operand))
                raise ValueError(f"Unsupported operation: {type(node)}")
    
            result = eval_node(tree.body)
            return f"{expression} = {result:.4f}"
        except Exception as e:
            return f"Calculation error: {e}"
    
    
    def dispatch_tool(name: str, args: dict) -> str:
        """Route a tool call from the agent to the right implementation."""
        if name == "search_docs":
            return search_docs(args["query"])
        if name == "get_entity":
            return get_entity(args["name"])
        if name == "calculate":
            return calculate(args["expression"])
        return f"Unknown tool: {name}"
    
    
    # ─────────────────────────────────────────────────────────────────
    #  The ReAct Agent Loop
    # ─────────────────────────────────────────────────────────────────
    
    def run_agent(question: str, max_iterations: int = 5) -> str:
        """
        Runs the Reason + Act (ReAct) loop.
    
        Each iteration:
          1. Send the conversation history to the LLM
          2. If the LLM calls a tool → execute it, append result, continue
          3. If the LLM gives a text response → that's the final answer, stop
    
        max_iterations prevents infinite loops if the agent gets stuck.
        """
        client, model = get_llm()
    
        # Conversation starts with a system prompt + the user's question
        messages = [
            {
                "role": "system",
                "content": (
                    "You are a research assistant with tools to look up space industry facts. "
                    "For each question, think about what information you need, use your tools to "
                    "retrieve it, and synthesize a clear, factual answer. "
                    "If a question requires multiple lookups or calculations, do them step by step."
                )
            },
            {"role": "user", "content": question}
        ]
    
        for iteration in range(max_iterations):
            response = client.chat.completions.create(
                model=model,
                messages=messages,
                tools=TOOLS,
                tool_choice="auto"  # let the LLM decide whether to call a tool
            )
    
            message = response.choices[0].message
            finish_reason = response.choices[0].finish_reason
    
            # ── Case 1: LLM wants to call a tool ─────────────────────
            if finish_reason == "tool_calls" and message.tool_calls:
                messages.append(message)  # add assistant message with tool call
    
                for tool_call in message.tool_calls:
                    tool_name = tool_call.function.name
                    tool_args = json.loads(tool_call.function.arguments)
    
                    print(f"   🔧 Tool call [{iteration+1}]: {tool_name}({tool_args})")
                    result = dispatch_tool(tool_name, tool_args)
                    print(f"      Result: {result[:120]}...")
    
                    # Append the tool result so the LLM sees it in the next iteration
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": result
                    })
    
            # ── Case 2: LLM has a final answer ────────────────────────
            else:
                return message.content
    
        return "Agent reached maximum iterations without a final answer."
    
    
    # ─────────────────────────────────────────────────────────────────
    #  Run example queries
    # ─────────────────────────────────────────────────────────────────
    
    questions = [
        # Single-hop: one tool call should be enough
        "What is SpaceX's approach to rocket reusability?",
    
        # Multi-hop: needs two entity lookups then synthesis
        "Compare SpaceX and NASA's plans for Mars exploration.",
    
        # Tool-chaining: retrieve + calculate
        "What is the ISS orbital speed in miles per hour?",
    ]
    
    for question in questions:
        print(f"❓ {question}")
        answer = run_agent(question)
        print(f"   💬 {answer}\n")
    

    Калі яго вжываць

    • Адказы, якім трэба калькуляцыя па колькае етапаў
    • Неодназначныя запыты, дзе план не вядомы з самага пачатку
    • Спалучэнне выкарыстоўвання дадзеных з вычысленнямі аб іншымі дзеяннямі
    • Эксперыментальныя далейшыя канверсаціі

    Выберыце лёгкіяшы шаблон, який падходзіць да формату запита: «naive» для чыстых дэманав, «hybrid» для пошуку текста ў прыемным режыме, «graphs» для стварэння сяродзецкіх звязкаў, а «agents» калі сам план выявлення інфармацыі должен быць дынамічным.