Monen asiakirjan kysymys–vastaus-agentit

import chromadb
from chromadb.utils import embedding_functions
import os

client = chromadb.PersistentClient(path='./doc_index')
ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key=os.getenv('OPENAI_API_KEY'),
    model_name='text-embedding-3-small'
)
collection = client.get_or_create_collection('documents', embedding_function=ef)

def index_document(doc_id, doc_path, doc_title):
    # Parse and chunk
    chunks = pdf_to_chunks(doc_path, chunk_size=800, overlap=150)
    for i, chunk in enumerate(chunks):
        chunk_id = f'{doc_id}_chunk_{i}'
        collection.add(
            ids=[chunk_id],
            documents=[chunk['text']],
            metadatas=[{
                'doc_id': doc_id,
                'title': doc_title,
                'page': chunk['page'],
                'source_file': doc_path
            }]
        )
    print(f'Indexed {len(chunks)} chunks from: {doc_title}')

def retrieve_relevant_chunks(question, n_results=8):
    results = collection.query(
        query_texts=[question],
        n_results=n_results,
        include=['documents', 'metadatas', 'distances']
    )

    chunks = []
    for i in range(len(results['documents'][0])):
        chunks.append({
            'text':      results['documents'][0][i],
            'metadata':  results['metadatas'][0][i],
            'distance':  results['distances'][0][i],
            'relevance': 1 - results['distances'][0][i]  # cosine similarity proxy
        })

    # Sort by relevance
    chunks.sort(key=lambda x: x['relevance'], reverse=True)
    return chunks

def format_chunks_for_prompt(chunks, max_chars=4000):
    sections = []
    used_chars = 0

    for i, chunk in enumerate(chunks, 1):
        meta = chunk['metadata']
        header = f"[Source {i}: {meta['title']}, page {meta.get('page', '?')}]"
        content = chunk['text'][:600]
        entry = f'{header}\n{content}'

        if used_chars + len(entry) > max_chars:
            break

        sections.append(entry)
        used_chars += len(entry)

    return '\n\n'.join(sections)

QA_PROMPT = '''Answer the question based on the provided document excerpts.
Cite sources as [Source N]. If sources conflict, mention both views.

{context}

Question: {question}
Answer:'''

def answer_question(question):
    chunks = retrieve_relevant_chunks(question, n_results=6)
    context = format_chunks_for_prompt(chunks)
    return llm_call(QA_PROMPT.format(context=context, question=question))

def cross_document_compare(question, doc_ids):
    # Retrieve best chunks per document
    per_doc_chunks = {}
    for doc_id in doc_ids:
        results = collection.query(
            query_texts=[question],
            n_results=3,
            where={'doc_id': {'$eq': doc_id}}  # filter by document
        )
        if results['documents'][0]:
            per_doc_chunks[doc_id] = results['documents'][0]

    # Format with document labels
    context_parts = []
    for doc_id, texts in per_doc_chunks.items():
        doc_label = f'Document {doc_id}'
        combined = ' '.join(texts[:2])[:600]
        context_parts.append(f'== {doc_label} ==\n{combined}')

    comparison_context = '\n\n'.join(context_parts)
    return llm_call(f'Compare these documents to answer: {question}\n\n{comparison_context}')

CONFLICT_PROMPT = '''You are analyzing multiple document sources.
Some may contain conflicting information.

For each factual claim you make:
1. Cite the source document
2. If another source contradicts it, explicitly note the conflict
3. Indicate which source you believe is more authoritative, if possible

Document excerpts:
{context}

Question: {question}

Answer (with conflict notes where applicable):'''

def answer_with_conflict_detection(question):
    chunks = retrieve_relevant_chunks(question, n_results=8)
    context = format_chunks_for_prompt(chunks)
    return llm_call(CONFLICT_PROMPT.format(
        context=context, question=question
    ))

MIN_RELEVANCE = 0.72  # cosine similarity threshold

def retrieve_above_threshold(question, n_results=10, threshold=MIN_RELEVANCE):
    chunks = retrieve_relevant_chunks(question, n_results=n_results)
    relevant = [c for c in chunks if c['relevance'] >= threshold]

    print(f'Retrieved: {len(chunks)}, Above threshold: {len(relevant)}')

    if not relevant:
        # Fallback: use top 3 even if below threshold
        return chunks[:3]

    return relevant

def answer_with_threshold(question):
    chunks = retrieve_above_threshold(question)
    if not chunks:
        return 'I could not find relevant information in the indexed documents.'
    context = format_chunks_for_prompt(chunks)
    return llm_call(QA_PROMPT.format(context=context, question=question))

import re

def extract_citations(answer_text, chunks):
    # Find all [Source N] references in the answer
    cited_nums = set(int(m) for m in re.findall(r'\[Source (\d+)\]', answer_text))

    citations = []
    for num in sorted(cited_nums):
        idx = num - 1
        if idx < len(chunks):
            meta = chunks[idx]['metadata']
            citations.append({
                'source_num': num,
                'title':     meta.get('title', 'Unknown'),
                'page':      meta.get('page', 'N/A'),
                'file':      meta.get('source_file', '')
            })

    return citations

def answer_with_citations(question):
    chunks = retrieve_above_threshold(question)
    context = format_chunks_for_prompt(chunks)
    answer = llm_call(QA_PROMPT.format(context=context, question=question))
    citations = extract_citations(answer, chunks)
    return {'answer': answer, 'citations': citations}

# pip install sentence-transformers
from sentence_transformers import CrossEncoder

reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

def rerank_chunks(question, chunks, top_k=4):
    # Score each chunk against the question
    pairs = [(question, c['text']) for c in chunks]
    scores = reranker.predict(pairs)

    # Attach scores and re-sort
    scored_chunks = list(zip(scores, chunks))
    scored_chunks.sort(key=lambda x: x[0], reverse=True)

    top_chunks = [chunk for _, chunk in scored_chunks[:top_k]]
    print(f'Re-ranked {len(chunks)} chunks -> kept top {top_k}')
    return top_chunks

def answer_with_reranking(question):
    # Retrieve more initially
    initial_chunks = retrieve_relevant_chunks(question, n_results=12)
    # Re-rank for precision
    top_chunks = rerank_chunks(question, initial_chunks, top_k=4)
    context = format_chunks_for_prompt(top_chunks)
    return llm_call(QA_PROMPT.format(context=context, question=question))

def retrieve_filtered(question, filters=None, n_results=8):
    query_kwargs = {
        'query_texts': [question],
        'n_results': n_results,
        'include': ['documents', 'metadatas', 'distances']
    }

    # ChromaDB metadata filters
    # Example: {'doc_id': 'contract_2024', 'year': {'$gte': 2023}}
    if filters:
        query_kwargs['where'] = filters

    results = collection.query(**query_kwargs)
    return [
        {'text': t, 'metadata': m, 'relevance': 1 - d}
        for t, m, d in zip(
            results['documents'][0],
            results['metadatas'][0],
            results['distances'][0]
        )
    ]

# Example usage
chunks = retrieve_filtered(
    'What are the payment terms?',
    filters={'doc_id': {'$in': ['contract_a', 'contract_b']}}
)

import os
import hashlib

# Track indexed documents by file hash
index_registry = {}  # {filepath: {hash, doc_id, indexed_at}}

def file_hash(filepath):
    with open(filepath, 'rb') as f:
        return hashlib.md5(f.read()).hexdigest()

def index_if_new_or_changed(filepath, title):
    fhash = file_hash(filepath)
    existing = index_registry.get(filepath)

    if existing and existing['hash'] == fhash:
        print(f'Skipping unchanged: {title}')
        return existing['doc_id']

    if existing:
        # Remove old chunks from vector store
        collection.delete(where={'doc_id': {'': existing['doc_id']}})
        print(f'Re-indexing updated: {title}')
    else:
        print(f'Indexing new: {title}')

    doc_id = hashlib.md5(filepath.encode()).hexdigest()[:8]
    index_document(doc_id, filepath, title)
    index_registry[filepath] = {'hash': fhash, 'doc_id': doc_id}
    return doc_id

def sync_document_directory(directory):
    pdf_files = [f for f in os.listdir(directory) if f.endswith('.pdf')]
    for fname in pdf_files:
        fpath = os.path.join(directory, fname)
        title = fname.replace('.pdf', '').replace('_', ' ').title()
        index_if_new_or_changed(fpath, title)
    print(f'Sync complete: {len(pdf_files)} files processed')