Объединение векторного и графового поиска

Зачем нужен гибридный поиск

Итоги векторного поиска

import openai
import numpy as np

client = openai.OpenAI(api_key='sk-...')

def embed(text: str) -> list:
    response = client.embeddings.create(
        model='text-embedding-3-small',
        input=text
    )
    return response.data[0].embedding

def cosine_similarity(a: list, b: list) -> float:
    a_arr = np.array(a)
    b_arr = np.array(b)
    return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))

# Simple in-memory vector store
class SimpleVectorStore:
    def __init__(self):
        self.documents = []
    
    def add(self, text: str, metadata: dict):
        embedding = embed(text)
        self.documents.append({'text': text, 'embedding': embedding, 'metadata': metadata})
    
    def search(self, query: str, top_k: int = 5) -> list:
        query_emb = embed(query)
        scored = [
            (cosine_similarity(query_emb, doc['embedding']), doc)
            for doc in self.documents
        ]
        scored.sort(key=lambda x: x[0], reverse=True)
        return [doc for _, doc in scored[:top_k]]

Итоги поиска по графу

from neo4j import GraphDatabase

driver = GraphDatabase.driver('bolt://localhost:7687', auth=('neo4j', 'password'))

def get_entity_context(entity_name: str) -> dict:
    with driver.session() as session:
        # Get node properties
        result = session.run(
            'MATCH (n {name: $name}) RETURN n, labels(n) AS labels LIMIT 1',
            name=entity_name
        )
        record = result.single()
        if not record:
            return {}
        
        node_data = dict(record['n'])
        node_labels = record['labels']
        
        # Get connected entities
        conn_result = session.run(
            'MATCH (n {name: $name})-[r]-(connected) '
            'RETURN type(r) AS rel_type, connected.name AS connected_name, labels(connected) AS connected_labels '
            'LIMIT 20',
            name=entity_name
        )
        connections = [dict(r) for r in conn_result]
        
        return {
            'name': entity_name,
            'labels': node_labels,
            'properties': node_data,
            'connections': connections
        }

Чередование результатов

def interleave_results(vector_results: list, graph_results: list) -> list:
    combined = []
    v_idx, g_idx = 0, 0
    
    while v_idx < len(vector_results) or g_idx < len(graph_results):
        if v_idx < len(vector_results):
            item = vector_results[v_idx]
            item['source'] = 'vector'
            combined.append(item)
            v_idx += 1
        
        if g_idx < len(graph_results):
            item = graph_results[g_idx]
            item['source'] = 'graph'
            combined.append(item)
            g_idx += 1
    
    return combined

# Example
vector_docs = [
    {'text': 'Alice led the machine learning initiative at Acme', 'score': 0.92},
    {'text': 'Machine learning best practices guide', 'score': 0.85},
]
graph_context = [
    {'name': 'Alice', 'type': 'Person', 'connections': ['Acme Corp', 'Bob']},
]

fused = interleave_results(vector_docs, graph_context)
for item in fused:
    print(f"[{item['source']}]", item.get('text') or item.get('name'))

Взвешенное объединение

def weighted_fusion(vector_results: list, graph_results: list, alpha: float = 0.6) -> list:
    '''
    alpha: weight for vector results (0.0 = pure graph, 1.0 = pure vector)
    '''
    all_results = []
    
    # Normalize vector scores (already in 0-1 range for cosine)
    for i, res in enumerate(vector_results):
        # Positional score: first result gets highest
        positional_score = 1.0 - (i / max(len(vector_results), 1))
        combined = alpha * res.get('score', positional_score)
        all_results.append({
            'content': res,
            'source': 'vector',
            'final_score': combined
        })
    
    # Graph results: score by relevance (e.g., connection count)
    for i, res in enumerate(graph_results):
        positional_score = 1.0 - (i / max(len(graph_results), 1))
        combined = (1 - alpha) * positional_score
        all_results.append({
            'content': res,
            'source': 'graph',
            'final_score': combined
        })
    
    # Sort by final score
    all_results.sort(key=lambda x: x['final_score'], reverse=True)
    return all_results

print('Weighted fusion function defined (alpha=0.6 favors vector)')

Слияние по обратным рангам

def reciprocal_rank_fusion(result_lists: list, k: int = 60) -> list:
    '''
    result_lists: list of lists, each containing dicts with an 'id' field
    k: constant to reduce impact of high rankings (typically 60)
    '''
    scores = {}
    all_items = {}
    
    for result_list in result_lists:
        for rank, item in enumerate(result_list):
            item_id = item.get('id') or item.get('text', '')[:50]
            if item_id not in scores:
                scores[item_id] = 0.0
                all_items[item_id] = item
            scores[item_id] += 1.0 / (k + rank + 1)
    
    sorted_ids = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)
    return [
        {**all_items[id_], 'rrf_score': scores[id_]}
        for id_ in sorted_ids
    ]

vector_list = [{'id': 'doc1', 'text': 'About Alice'}, {'id': 'doc3', 'text': 'About AI'}]
graph_list = [{'id': 'doc2', 'text': 'Alice connections'}, {'id': 'doc1', 'text': 'About Alice'}]

fused = reciprocal_rank_fusion([vector_list, graph_list])
for item in fused:
    print(f"{item['id']}: RRF score {item['rrf_score']:.4f}")

Гибридный поиск с привязкой к сущностям

import spacy

nlp = spacy.load('en_core_web_sm')

def entity_anchored_retrieval(query: str, vector_store, graph_driver) -> dict:
    # Step 1: Extract entities from query
    doc = nlp(query)
    entities = [ent.text for ent in doc.ents if ent.label_ in ['PERSON', 'ORG', 'GPE']]
    
    # Step 2: Get graph context for entities
    graph_contexts = {}
    for entity in entities:
        context = get_entity_context(entity)
        if context:
            graph_contexts[entity] = context
    
    # Step 3: Enrich query with graph context
    enriched_query = query
    if graph_contexts:
        context_str = ' '.join([
            f"{name} works at {', '.join([c['connected_name'] for c in ctx.get('connections', [])[:3]])}"
            for name, ctx in graph_contexts.items()
        ])
        enriched_query = f'{query} Context: {context_str}'
    
    # Step 4: Vector search with enriched query
    vector_results = vector_store.search(enriched_query, top_k=5)
    
    return {
        'entities_found': entities,
        'graph_contexts': graph_contexts,
        'vector_results': vector_results
    }

Создание пакета контекста

def build_context_package(vector_results: list, graph_contexts: dict, max_tokens: int = 3000) -> str:
    sections = []
    
    # Graph entity context section
    if graph_contexts:
        graph_section = ['## Entity Context from Knowledge Graph']
        for entity_name, context in graph_contexts.items():
            connections = context.get('connections', [])
            conn_summary = ', '.join([
                f"{c['connected_name']} ({c['rel_type']})"
                for c in connections[:5]
            ])
            graph_section.append(f'**{entity_name}**: connected to {conn_summary}')
        sections.append('\n'.join(graph_section))
    
    # Vector search results section
    if vector_results:
        vector_section = ['## Relevant Documents']
        for i, doc in enumerate(vector_results[:5]):
            text = doc.get('text', '')[:500]  # Truncate long docs
            vector_section.append(f'{i+1}. {text}')
        sections.append('\n'.join(vector_section))
    
    context_package = '\n\n'.join(sections)
    # Rough token estimate (1 token ~ 4 chars)
    if len(context_package) > max_tokens * 4:
        context_package = context_package[:max_tokens * 4]
    
    return context_package

if __name__ == '__main__':
    demo_vector = [{'text': 'Refunds are processed within 5 business days of approval.'}]
    demo_graph = {'Acme Corp': {'connections': [{'connected_name': 'Jane Doe', 'rel_type': 'employs'}]}}
    print(build_context_package(demo_vector, demo_graph))

Асинхронный параллельный поиск

import asyncio
from concurrent.futures import ThreadPoolExecutor

executor = ThreadPoolExecutor(max_workers=4)

async def async_vector_search(query: str, vector_store) -> list:
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(executor, vector_store.search, query, 5)

async def async_graph_lookup(entities: list) -> dict:
    loop = asyncio.get_event_loop()
    results = {}
    for entity in entities:
        context = await loop.run_in_executor(executor, get_entity_context, entity)
        if context:
            results[entity] = context
    return results

async def hybrid_retrieval_async(query: str, entities: list, vector_store) -> dict:
    # Run vector search and graph lookup in parallel
    vector_task = async_vector_search(query, vector_store)
    graph_task = async_graph_lookup(entities)
    
    vector_results, graph_contexts = await asyncio.gather(vector_task, graph_task)
    
    return {
        'vector': vector_results,
        'graph': graph_contexts
    }

print('Async parallel retrieval functions defined')

Кэширование результатов поиска

import hashlib
import time

class HybridRetrievalCache:
    def __init__(self, vector_ttl: int = 300, graph_ttl: int = 600):
        self.vector_cache = {}
        self.graph_cache = {}
        self.vector_ttl = vector_ttl
        self.graph_ttl = graph_ttl
    
    def _key(self, value: str) -> str:
        return hashlib.md5(value.encode()).hexdigest()[:12]
    
    def get_vector(self, query: str):
        k = self._key(query)
        entry = self.vector_cache.get(k)
        if entry and time.time() - entry['ts'] < self.vector_ttl:
            return entry['data']
        return None
    
    def set_vector(self, query: str, results: list):
        self.vector_cache[self._key(query)] = {'data': results, 'ts': time.time()}
    
    def get_graph(self, entity: str):
        k = self._key(entity)
        entry = self.graph_cache.get(k)
        if entry and time.time() - entry['ts'] < self.graph_ttl:
            return entry['data']
        return None
    
    def set_graph(self, entity: str, context: dict):
        self.graph_cache[self._key(entity)] = {'data': context, 'ts': time.time()}

cache = HybridRetrievalCache()
print('Hybrid retrieval cache initialized')

Выбор весов поиска

def auto_tune_alpha(query: str) -> float:
    query_lower = query.lower()
    
    # High graph weight for relational questions
    relational_keywords = [
        'who', 'founded', 'works at', 'connected to',
        'related to', 'partner', 'owns', 'acquired'
    ]
    
    # High vector weight for content questions
    content_keywords = [
        'explain', 'describe', 'what is', 'how does',
        'tell me about', 'documents about', 'find information'
    ]
    
    relational_count = sum(1 for kw in relational_keywords if kw in query_lower)
    content_count = sum(1 for kw in content_keywords if kw in query_lower)
    
    if relational_count > content_count:
        return 0.3  # Graph-heavy
    elif content_count > relational_count:
        return 0.7  # Vector-heavy
    else:
        return 0.5  # Balanced

queries = [
    'Who founded Tesla?',
    'Explain transformer architecture',
    'What companies is Elon Musk connected to?'
]
for q in queries:
    print(f'alpha={auto_tune_alpha(q):.1f} for: {q}')

Проверка знаний: гибридный поиск

Итоги гибридного поиска