Объединение веб-поиска с RAG

class HybridRetrievalAgent:
    def __init__(self, vector_store, search_client):
        self.vector_store = vector_store  # e.g., ChromaDB or FAISS
        self.search_client = search_client  # e.g., TavilyClient

    def answer(self, question):
        route = self.classify_question(question)

        if route == 'static':
            context = self.vector_store.query(question, n_results=5)
        elif route == 'current':
            context = self.web_search(question)
        else:  # 'both'
            local = self.vector_store.query(question, n_results=3)
            web = self.web_search(question)
            context = local + web

        return self.generate_answer(question, context)

if __name__ == '__main__':
    class DemoAgent(HybridRetrievalAgent):
        def classify_question(self, question):
            return 'static' if 'company' in question.lower() else 'current'
        def web_search(self, question):
            return [('Web result about ' + question, {})]
        def generate_answer(self, question, context):
            return f'Answer using {len(context)} context item(s).'

    class FakeVectorStore:
        def query(self, question, n_results=5):
            return [('Local doc snippet', {})]

    agent = DemoAgent(FakeVectorStore(), search_client=None)
    print(agent.answer('What is our company policy on refunds?'))

ROUTING_PROMPT = '''Classify this question into one of three categories:
- "static": answered from company documents, product docs, or stable technical knowledge
- "current": requires up-to-date information (news, prices, recent events, latest releases)
- "both": needs both company context and current information

Question: {question}

Respond with exactly one word: static, current, or both.'''

def classify_question(question):
    response = llm_call(ROUTING_PROMPT.format(question=question))
    route = response.strip().lower()
    if route not in ('static', 'current', 'both'):
        return 'both'  # safe default
    return route

import chromadb
from chromadb.utils import embedding_functions
import os

client = chromadb.PersistentClient(path='./vector_db')

ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key=os.getenv('OPENAI_API_KEY'),
    model_name='text-embedding-3-small'
)

collection = client.get_or_create_collection(
    name='company_docs',
    embedding_function=ef
)

def index_document(doc_id, text, metadata=None):
    collection.add(
        ids=[doc_id],
        documents=[text],
        metadatas=[metadata or {}]
    )

def local_retrieve(question, n_results=5):
    results = collection.query(
        query_texts=[question],
        n_results=n_results
    )
    return list(zip(results['documents'][0], results['metadatas'][0]))

from tavily import TavilyClient
import os

tavily = TavilyClient(api_key=os.getenv('TAVILY_API_KEY'))

def web_retrieve(question, n_results=3):
    results = tavily.search(
        query=question,
        max_results=n_results,
        search_depth='basic'
    )
    # Normalize to same format as local results
    return [
        (
            r['content'][:600],  # text
            {'source': r['url'], 'title': r['title'], 'type': 'web'}  # metadata
        )
        for r in results.get('results', [])
    ]

def merge_results(local_results, web_results):
    merged = []

    for text, meta in local_results:
        merged.append({
            'content': text,
            'source': meta.get('source', 'internal document'),
            'type': 'local',
            'title': meta.get('title', 'Company Document')
        })

    for text, meta in web_results:
        merged.append({
            'content': text,
            'source': meta.get('source', 'web'),
            'type': 'web',
            'title': meta.get('title', 'Web Result')
        })

    return merged

def format_merged_for_prompt(merged_results):
    parts = []
    for i, r in enumerate(merged_results, 1):
        tag = '[INTERNAL]' if r['type'] == 'local' else '[WEB]'
        parts.append(f'[{i}] {tag} {r["title"]}\n{r["content"]}')
    return '\n\n'.join(parts)

if __name__ == '__main__':
    local = [('Refunds are processed within 5 business days.', {'source': 'handbook', 'title': 'Refund Policy'})]
    web = [('Company X reported Q2 earnings today.', {'source': 'reuters.com', 'title': 'Q2 Earnings'})]
    merged = merge_results(local, web)
    print(format_merged_for_prompt(merged))

CURRENT_EVENTS_SIGNALS = [
    'latest', 'current', 'today', 'now', 'recent',
    'this week', 'this month', 'this year',
    'just released', 'new version', 'updated',
    'price', 'stock', 'news', 'announcement',
    '2024', '2025'
]

STATIC_SIGNALS = [
    'how does', 'what is', 'explain', 'tutorial',
    'documentation', 'our product', 'company policy',
    'internal', 'handbook'
]

def fast_route(question):
    lower = question.lower()
    current_score = sum(1 for s in CURRENT_EVENTS_SIGNALS if s in lower)
    static_score = sum(1 for s in STATIC_SIGNALS if s in lower)

    if current_score > static_score:
        return 'current'
    elif static_score > current_score:
        return 'static'
    else:
        return 'both'

if __name__ == '__main__':
    for q in ['What is our company handbook policy on PTO?', 'What is the latest stock price today?']:
        print(f'{fast_route(q)!r} <- "{q}"')

HYBRID_ANSWER_PROMPT = '''You are answering a question using two types of sources:
- [INTERNAL] sources: company documents (may be outdated)
- [WEB] sources: current web information

For factual claims about current state (prices, versions, availability):
  PREFER [WEB] sources over [INTERNAL] ones.
For company-specific processes, policies, and architecture:
  PREFER [INTERNAL] sources.

If sources conflict, note the discrepancy in your answer.

Sources:
{sources}

Question: {question}
Answer:'''

def generate_hybrid_answer(question, merged_results):
    sources_text = format_merged_for_prompt(merged_results)
    return llm_call(HYBRID_ANSWER_PROMPT.format(
        sources=sources_text,
        question=question
    ))

from datetime import datetime, timedelta

STALENESS_THRESHOLD_DAYS = 90

def check_staleness(metadata):
    indexed_at = metadata.get('indexed_at')
    if not indexed_at:
        return False  # unknown age — assume fresh
    indexed_date = datetime.fromisoformat(indexed_at)
    age = datetime.now() - indexed_date
    return age > timedelta(days=STALENESS_THRESHOLD_DAYS)

def smart_retrieve(question, route):
    local_results = []
    web_results = []

    if route in ('static', 'both'):
        local_results = local_retrieve(question, n_results=4)
        # Check if any local results are stale
        stale = any(check_staleness(meta) for _, meta in local_results)
        if stale:
            print('Stale local docs — adding web search')
            web_results = web_retrieve(question, n_results=2)

    if route in ('current', 'both'):
        web_results = web_retrieve(question, n_results=3)

    return merge_results(local_results, web_results)

def score_result(result, query_embedding):
    score = 0.5  # base score

    # Recency bonus for web results
    if result.get('type') == 'web':
        pub_date = result.get('published_date', '')
        if '2024' in pub_date or '2025' in pub_date:
            score += 0.2

    # Embedding similarity to query
    if result.get('content'):
        result_emb = embed(result['content'][:500])
        sim = cosine_similarity(query_embedding, result_emb)
        score += sim * 0.3

    # Domain authority
    from urllib.parse import urlparse
    domain = urlparse(result.get('source', '')).netloc
    if any(auth in domain for auth in ['docs.', 'developer.', 'official.']):
        score += 0.1

    return min(score, 1.0)

def hybrid_answer(question):
    # 1. Route (fast heuristic first, LLM fallback for ambiguous)
    route = fast_route(question)
    if route == 'both':
        route = classify_question(question)  # LLM for ambiguous cases

    print(f'Route: {route}')

    # 2. Retrieve
    merged = smart_retrieve(question, route)

    if not merged:
        return 'I could not find relevant information to answer your question.'

    # 3. Generate
    answer = generate_hybrid_answer(question, merged)
    return answer

# Usage
print(hybrid_answer('What is our refund policy?'))   # -> static/local
print(hybrid_answer('What is GPT-4 pricing today?')) # -> current/web