Агент решает: нужен ли поиск, по каким документам

import re
import json
from enum import StrEnum
from openai import OpenAI

client = OpenAI()


class Intent(StrEnum):
    DIRECT = "direct"
    SEARCH = "search"


# Эвристики для быстрой фильтрации очевидных случаев
DIRECT_PATTERNS = [
    r'^(привет|здравствуй|добрый|добрый день|hi|hello)',
    r'^\s*\d+[\s]*[\+\-\*\/\^][\s]*\d+',           # Математика
    r'^(переведи|перефразируй|исправь|объясни\s+слово)',
    r'^(спасибо|ок|понял|ясно|хорошо)\s*[\.!]?\s*$',
]

INTENT_PROMPT = """Classify this user query. Be precise.

DIRECT — answer from general knowledge, NO company docs needed:
  - Greetings, small talk
  - Basic math or unit conversions
  - Generic CS/programming definitions
  - Language/translation tasks

SEARCH — needs company-specific knowledge base:
  - Product features, pricing, integrations
  - Company policies (HR, legal, ops)
  - Specific versions, release notes, changelogs
  - Support questions about this product

Query: {query}

Return JSON: {{"intent": "DIRECT" | "SEARCH", "reason": "one line"}}"""


def classify_intent(query: str, use_llm_fallback: bool = True) -> Intent:
    # 1. Быстрые правила
    q_lower = query.strip().lower()
    if any(re.match(p, q_lower) for p in DIRECT_PATTERNS):
        return Intent.DIRECT

    if not use_llm_fallback:
        return Intent.SEARCH  # Безопасный дефолт

    # 2. LLM-классификация для неоднозначных случаев
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": INTENT_PROMPT.format(query=query)}],
        response_format={"type": "json_object"},
        temperature=0,
        max_tokens=60,
    )
    result = json.loads(resp.choices[0].message.content)
    return Intent(result["intent"].lower())


# Проверка
print(classify_intent("Привет!"))                       # Intent.DIRECT
print(classify_intent("Что такое async/await?"))        # Intent.DIRECT (общее знание)
print(classify_intent("Поддерживаете ли вы Kafka?"))   # Intent.SEARCH
print(classify_intent("Что нового в версии 2.1?"))     # Intent.SEARCH

import numpy as np
from sentence_transformers import SentenceTransformer
from dataclasses import dataclass

embed_model = SentenceTransformer("intfloat/multilingual-e5-large")


@dataclass
class Collection:
    name: str
    description: str           # Для embedding-роутинга
    keywords: list[str]        # Для keyword-роутинга
    embedding: np.ndarray | None = None


# Определяем коллекции с описаниями
COLLECTIONS: list[Collection] = [
    Collection(
        name="technical_docs",
        description="API documentation, SDK integration, webhooks, authentication, rate limits, changelog",
        keywords=["api", "sdk", "endpoint", "интеграция", "токен", "webhook", "ошибка", "код"],
    ),
    Collection(
        name="hr_policies",
        description="vacation days, sick leave, salary, benefits, employment rules, onboarding",
        keywords=["отпуск", "зарплата", "больничный", "льготы", "офис", "испытательный", "увольнение"],
    ),
    Collection(
        name="product_faq",
        description="product features, pricing plans, trial, comparison with competitors, enterprise",
        keywords=["цена", "тариф", "enterprise", "пробный", "функции", "сравнение", "конкурент"],
    ),
]


def _precompute_embeddings():
    """Считаем embeddings описаний один раз при старте."""
    for coll in COLLECTIONS:
        coll.embedding = embed_model.encode(
            coll.description, normalize_embeddings=True
        )


_precompute_embeddings()


def route_by_embedding(query: str, top_k: int = 1) -> list[str]:
    """Embedding-based роутинг: O(n) сравнений, без LLM."""
    q_emb = embed_model.encode(query, normalize_embeddings=True)
    scores = [
        (coll.name, float(np.dot(q_emb, coll.embedding)))
        for coll in COLLECTIONS
    ]
    scores.sort(key=lambda x: x[1], reverse=True)
    return [name for name, _ in scores[:top_k]]


ROUTING_PROMPT = """You have these knowledge bases:
{collections}

Which knowledge bases are needed to answer this query?
Select ALL that apply. If the query spans multiple topics, select multiple.

Query: {query}

Return JSON: {{"collections": ["name1", "name2"], "reason": "brief"}}"""


def route_by_llm(query: str) -> list[str]:
    """LLM-based роутинг: точнее, медленнее (+100-200ms)."""
    coll_desc = "\n".join(
        f"- {c.name}: {c.description}"
        for c in COLLECTIONS
    )
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": ROUTING_PROMPT.format(
            collections=coll_desc, query=query
        )}],
        response_format={"type": "json_object"},
        temperature=0,
        max_tokens=80,
    )
    result = json.loads(resp.choices[0].message.content)
    # Валидируем: возвращаем только существующие коллекции
    valid = {c.name for c in COLLECTIONS}
    return [c for c in result.get("collections", []) if c in valid]


def route_query(query: str, use_llm: bool = True) -> list[str]:
    """Гибридный роутинг: embedding для скорости + LLM для точности."""
    if use_llm:
        return route_by_llm(query)
    return route_by_embedding(query, top_k=1)


# Примеры
print(route_query("Как подключить OAuth?"))     # ['technical_docs']
print(route_query("Отпуск и тарифы — расскажи"))  # ['hr_policies', 'product_faq']
print(route_query("Что нового в v2.1 API?"))    # ['technical_docs']

import asyncio
from typing import Any

# Определение инструментов для OpenAI API
RAG_TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "search_technical_docs",
            "description": (
                "Search technical documentation: API reference, SDK integration guides, "
                "webhooks, authentication, rate limits, error codes, changelogs. "
                "Use for: 'how to integrate', 'API endpoint', 'SDK error', version info."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "Search query optimized for technical documentation",
                    }
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "search_hr_policies",
            "description": (
                "Search HR policies and company rules: vacation days, sick leave, "
                "salary, benefits, office hours, remote work, onboarding, termination. "
                "Use for: 'how many days', 'can I', 'what is the policy on'."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "HR policy search query"}
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "search_product_faq",
            "description": (
                "Search product FAQ: pricing plans, features, trial period, "
                "comparison with competitors, enterprise options, support tiers. "
                "Use for: 'how much', 'pricing', 'does it support', 'compared to'."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Product FAQ search query"}
                },
                "required": ["query"],
            },
        },
    },
]


async def execute_tool(tool_name: str, arguments: dict, collections: dict) -> str:
    """Выполняет поиск в нужной коллекции."""
    collection_map = {
        "search_technical_docs": "technical_docs",
        "search_hr_policies": "hr_policies",
        "search_product_faq": "product_faq",
    }
    coll_name = collection_map[tool_name]
    db = collections[coll_name]
    results = db.search(arguments["query"], top_k=4)
    return "\n\n---\n\n".join(r.page_content for r in results)


SYSTEM_PROMPT = """You are a helpful assistant for our company.
Answer questions using the provided tools to search our knowledge bases.
If a question doesn't require searching (greetings, general knowledge, math),
answer directly without using any tools.
Always answer in the same language as the question."""


async def agentic_rag(question: str, collections: dict) -> str:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": question},
    ]

    # Шаг 1: Первый вызов — агент решает, нужны ли инструменты
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        tools=RAG_TOOLS,
        tool_choice="auto",  # Агент сам решает
        temperature=0.1,
    )
    msg = response.choices[0].message

    # Если инструментов нет — прямой ответ (грит, математика, общее знание)
    if not msg.tool_calls:
        return msg.content

    # Шаг 2: Выполняем все инструменты параллельно
    messages.append(msg.model_dump(exclude_unset=True))

    tool_results = await asyncio.gather(*[
        execute_tool(tc.function.name, json.loads(tc.function.arguments), collections)
        for tc in msg.tool_calls
    ])

    for tc, result in zip(msg.tool_calls, tool_results):
        messages.append({
            "role": "tool",
            "tool_call_id": tc.id,
            "content": result,
        })

    # Шаг 3: Финальный ответ с контекстом
    final = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        temperature=0.1,
    )
    return final.choices[0].message.content

REWRITE_PROMPT = """Rewrite the user's question as a search query optimized for this knowledge base.

Knowledge base: {collection_name}
Description: {collection_desc}

Rules:
- Use technical/domain-specific terminology from the knowledge base
- Remove conversational filler ("do you have", "is there", "can I")
- Focus on the core concept being asked about
- Keep it concise (5-15 words)

User question: {question}
Search query:"""


@dataclass
class CollectionQuery:
    collection: str
    query: str


def rewrite_for_collection(question: str, collection_name: str) -> str:
    """Переформулирует вопрос под конкретную коллекцию."""
    coll = next(c for c in COLLECTIONS if c.name == collection_name)
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": REWRITE_PROMPT.format(
            collection_name=collection_name,
            collection_desc=coll.description,
            question=question,
        )}],
        temperature=0,
        max_tokens=40,
    )
    return resp.choices[0].message.content.strip()


async def multi_collection_search(
    question: str,
    target_collections: list[str],
    dbs: dict,
) -> list[str]:
    """Параллельный поиск с адаптацией запроса под каждую коллекцию."""
    async def search_one(coll_name: str) -> list[str]:
        # Переформулируем под коллекцию
        query = rewrite_for_collection(question, coll_name)
        results = dbs[coll_name].search(query, top_k=3)
        return [f"[{coll_name}] {r.page_content}" for r in results]

    all_results = await asyncio.gather(*[
        search_one(coll) for coll in target_collections
    ])
    # Плоский список чанков из всех коллекций
    return [chunk for results in all_results for chunk in results]


# Пример
# question = "Есть ли нормальная поддержка платёжек?"
# rewrite_for_collection(question, "technical_docs")
# → "payment gateway integration PCI DSS compliance"

from dataclasses import dataclass, field
from typing import Protocol


class VectorDB(Protocol):
    def search(self, query: str, top_k: int = 4) -> list: ...


@dataclass
class AgentConfig:
    model: str = "gpt-4o-mini"
    router_model: str = "gpt-4o-mini"  # Можно дешевле для роутинга
    top_k: int = 4
    temperature: float = 0.1
    use_llm_routing: bool = True


@dataclass
class AgenticRAG:
    dbs: dict[str, VectorDB]
    config: AgentConfig = field(default_factory=AgentConfig)

    SYSTEM_PROMPT = """You are a helpful assistant. Answer questions based on the provided context.
If context is provided, use it. Cite specific information when possible.
If no context is provided, answer from general knowledge.
Answer in the same language as the question."""

    async def answer(self, question: str) -> dict:
        """Полный цикл: classify → route → search → generate."""
        trace = {"question": question}

        # ── Шаг 1: Нужен ли поиск? ──────────────────────
        intent = classify_intent(question)
        trace["intent"] = intent

        if intent == Intent.DIRECT:
            # Прямой ответ без retrieval
            resp = client.chat.completions.create(
                model=self.config.model,
                messages=[
                    {"role": "system", "content": self.SYSTEM_PROMPT},
                    {"role": "user", "content": question},
                ],
                temperature=self.config.temperature,
            )
            return {"answer": resp.choices[0].message.content, **trace, "contexts": []}

        # ── Шаг 2: Роутинг коллекции ─────────────────────
        target_colls = route_query(question, use_llm=self.config.use_llm_routing)
        trace["collections"] = target_colls

        # Оставляем только существующие коллекции
        target_colls = [c for c in target_colls if c in self.dbs]
        if not target_colls:
            target_colls = list(self.dbs.keys())[:1]  # Fallback

        # ── Шаг 3: Параллельный поиск ────────────────────
        contexts = await multi_collection_search(question, target_colls, self.dbs)
        trace["context_chunks"] = len(contexts)

        # ── Шаг 4: Генерация ответа ───────────────────────
        context_text = "\n\n---\n\n".join(contexts[:self.config.top_k * len(target_colls)])
        resp = client.chat.completions.create(
            model=self.config.model,
            messages=[
                {"role": "system", "content": self.SYSTEM_PROMPT},
                {"role": "user", "content": (
                    f"Context:\n{context_text}\n\nQuestion: {question}"
                )},
            ],
            temperature=self.config.temperature,
        )
        return {
            "answer": resp.choices[0].message.content,
            **trace,
            "contexts": contexts,
        }


# Использование
async def main():
    rag = AgenticRAG(
        dbs={
            "technical_docs": technical_db,
            "hr_policies": hr_db,
            "product_faq": faq_db,
        },
        config=AgentConfig(top_k=4, use_llm_routing=True),
    )

    questions = [
        "Привет! Как дела?",                   # DIRECT → 0 DB calls
        "Сколько дней отпуска в год?",          # hr_policies
        "Как настроить OAuth callback URL?",    # technical_docs
        "Цены на Enterprise и SLA гарантии?",  # product_faq
    ]

    for q in questions:
        result = await rag.answer(q)
        print(f"Q: {q}")
        print(f"  Intent: {result['intent']}, Collections: {result.get('collections', [])}")
        print(f"  Answer: {result['answer'][:100]}...")
        print()

# Самый простой Agentic RAG: один инструмент, агент решает сам
from openai import OpenAI

client = OpenAI()

def search_docs(query: str) -> str:
    # ваш vector search здесь
    results = vector_db.search(query, top_k=4)
    return "\n---\n".join(r.page_content for r in results)

tools = [{
    "type": "function",
    "function": {
        "name": "search_docs",
        "description": "Search company knowledge base. Use for company-specific questions.",
        "parameters": {
            "type": "object",
            "properties": {"query": {"type": "string"}},
            "required": ["query"],
        },
    },
}]

def answer(question: str) -> str:
    messages = [
        {"role": "system", "content": "Answer using tools when needed. Direct answer for greetings/math."},
        {"role": "user", "content": question},
    ]
    resp = client.chat.completions.create(
        model="gpt-4o-mini", messages=messages, tools=tools, tool_choice="auto"
    )
    msg = resp.choices[0].message
    if not msg.tool_calls:
        return msg.content  # Прямой ответ

    # Выполняем поиск
    context = search_docs(json.loads(msg.tool_calls[0].function.arguments)["query"])
    messages.extend([
        msg.model_dump(exclude_unset=True),
        {"role": "tool", "tool_call_id": msg.tool_calls[0].id, "content": context},
    ])
    return client.chat.completions.create(
        model="gpt-4o-mini", messages=messages
    ).choices[0].message.content