index_documents.py

import os, sys, time, hashlib, argparse
from pathlib import Path
from dotenv import load_dotenv
import tiktoken
from openai import AzureOpenAI
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex, SearchField, SearchFieldDataType, SimpleField,
    SearchableField, VectorSearch, HnswAlgorithmConfiguration,
    VectorSearchProfile, SemanticConfiguration, SemanticSearch,
    SemanticPrioritizedFields, SemanticField,
)
from azure.core.credentials import AzureKeyCredential
from azure.storage.blob import BlobServiceClient

load_dotenv()
OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
EMBEDDING_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT", "embedding-small")
API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION", "2024-10-21")
SEARCH_ENDPOINT = os.getenv("AZURE_SEARCH_ENDPOINT")
SEARCH_API_KEY = os.getenv("AZURE_SEARCH_API_KEY")
INDEX_NAME = os.getenv("AZURE_SEARCH_INDEX_NAME", "dispute-knowledge-base")
STORAGE_CONNECTION_STRING = os.getenv("AZURE_STORAGE_CONNECTION_STRING")
STORAGE_CONTAINER = os.getenv("AZURE_STORAGE_CONTAINER", "dispute-documents")
CHUNK_SIZE = 500
CHUNK_OVERLAP = 100
EMBEDDING_DIMENSIONS = 1536

def read_local_documents(folder_path):
    documents = []
    folder = Path(folder_path)
    for file_path in sorted(folder.glob("*.txt")):
        print(f"  Reading: {file_path.name}")
        content = file_path.read_text(encoding="utf-8")
        documents.append({"filename": file_path.name, "content": content})
    print(f"  Total documents read: {len(documents)}")
    return documents

def read_blob_documents():
    documents = []
    blob_service = BlobServiceClient.from_connection_string(STORAGE_CONNECTION_STRING)
    container_client = blob_service.get_container_client(STORAGE_CONTAINER)
    for blob in container_client.list_blobs():
        if blob.name.endswith(".txt"):
            print(f"  Downloading: {blob.name}")
            blob_client = container_client.get_blob_client(blob.name)
            content = blob_client.download_blob().readall().decode("utf-8")
            documents.append({"filename": blob.name, "content": content})
    print(f"  Total documents read: {len(documents)}")
    return documents

def generate_chunk_id(filename, chunk_index):
    raw = f"{filename}_{chunk_index}"
    return hashlib.md5(raw.encode()).hexdigest()

def get_overlap_text(text, overlap_tokens, tokenizer):
    tokens = tokenizer.encode(text)
    if len(tokens) <= overlap_tokens:
        return text
    return tokenizer.decode(tokens[-overlap_tokens:])

def chunk_document(document):
    tokenizer = tiktoken.get_encoding("cl100k_base")
    content = document["content"]
    filename = document["filename"]
    paragraphs = content.split("\n")
    chunks = []
    current_chunk_text = ""
    current_chunk_tokens = 0
    chunk_index = 0
    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
        para_tokens = len(tokenizer.encode(para))
        if para_tokens > CHUNK_SIZE:
            sentences = para.replace(". ", ".\n").split("\n")
            for sentence in sentences:
                sent_tokens = len(tokenizer.encode(sentence))
                if current_chunk_tokens + sent_tokens > CHUNK_SIZE and current_chunk_text:
                    chunk_id = generate_chunk_id(filename, chunk_index)
                    chunks.append({"id": chunk_id, "content": current_chunk_text.strip(), "filename": filename, "chunk_index": chunk_index, "token_count": current_chunk_tokens})
                    chunk_index += 1
                    overlap_text = get_overlap_text(current_chunk_text, CHUNK_OVERLAP, tokenizer)
                    current_chunk_text = overlap_text + " " + sentence
                    current_chunk_tokens = len(tokenizer.encode(current_chunk_text))
                else:
                    current_chunk_text += " " + sentence
                    current_chunk_tokens += sent_tokens
        elif current_chunk_tokens + para_tokens > CHUNK_SIZE and current_chunk_text:
            chunk_id = generate_chunk_id(filename, chunk_index)
            chunks.append({"id": chunk_id, "content": current_chunk_text.strip(), "filename": filename, "chunk_index": chunk_index, "token_count": current_chunk_tokens})
            chunk_index += 1
            overlap_text = get_overlap_text(current_chunk_text, CHUNK_OVERLAP, tokenizer)
            current_chunk_text = overlap_text + "\n" + para
            current_chunk_tokens = len(tokenizer.encode(current_chunk_text))
        else:
            current_chunk_text += "\n" + para
            current_chunk_tokens += para_tokens
    if current_chunk_text.strip():
        chunk_id = generate_chunk_id(filename, chunk_index)
        chunks.append({"id": chunk_id, "content": current_chunk_text.strip(), "filename": filename, "chunk_index": chunk_index, "token_count": current_chunk_tokens})
    return chunks

def extract_document_title(filename):
    name = filename.replace(".txt", "")
    parts = name.split("_", 1)
    if len(parts) > 1 and parts[0].isdigit():
        name = parts[1]
    return name.replace("_", " ").title()

def generate_embeddings(chunks):
    client = AzureOpenAI(azure_endpoint=OPENAI_ENDPOINT, api_key=OPENAI_API_KEY, api_version=API_VERSION)
    total = len(chunks)
    batch_size = 16
    for i in range(0, total, batch_size):
        batch = chunks[i:i + batch_size]
        texts = [chunk["content"] for chunk in batch]
        print(f"  Embedding batch {i // batch_size + 1}/{(total + batch_size - 1) // batch_size} ({len(texts)} chunks)")
        response = client.embeddings.create(input=texts, model=EMBEDDING_DEPLOYMENT)
        for j, embedding_data in enumerate(response.data):
            chunks[i + j]["embedding"] = embedding_data.embedding
        if i + batch_size < total:
            time.sleep(0.5)
    return chunks

def create_search_index():
    index_client = SearchIndexClient(endpoint=SEARCH_ENDPOINT, credential=AzureKeyCredential(SEARCH_API_KEY))
    fields = [
        SimpleField(name="id", type=SearchFieldDataType.String, key=True, filterable=True),
        SearchableField(name="content", type=SearchFieldDataType.String, analyzer_name="en.microsoft"),
        SearchableField(name="filename", type=SearchFieldDataType.String, filterable=True, facetable=True),
        SearchableField(name="title", type=SearchFieldDataType.String, filterable=True),
        SimpleField(name="chunk_index", type=SearchFieldDataType.Int32, filterable=True, sortable=True),
        SimpleField(name="token_count", type=SearchFieldDataType.Int32, filterable=True),
        SearchField(name="embedding", type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
                    searchable=True, vector_search_dimensions=EMBEDDING_DIMENSIONS, vector_search_profile_name="vector-profile"),
    ]
    vector_search = VectorSearch(
        algorithms=[HnswAlgorithmConfiguration(name="hnsw-config")],
        profiles=[VectorSearchProfile(name="vector-profile", algorithm_configuration_name="hnsw-config")]
    )
    semantic_config = SemanticConfiguration(
        name="semantic-config",
        prioritized_fields=SemanticPrioritizedFields(
            content_fields=[SemanticField(field_name="content")],
            title_field=SemanticField(field_name="title")
        )
    )
    semantic_search = SemanticSearch(configurations=[semantic_config])
    index = SearchIndex(name=INDEX_NAME, fields=fields, vector_search=vector_search, semantic_search=semantic_search)
    try:
        index_client.delete_index(INDEX_NAME)
        print(f"  Deleted existing index: {INDEX_NAME}")
    except Exception:
        pass
    result = index_client.create_index(index)
    print(f"  Created index: {result.name}")
    return result

def upload_to_index(chunks):
    search_client = SearchClient(endpoint=SEARCH_ENDPOINT, index_name=INDEX_NAME, credential=AzureKeyCredential(SEARCH_API_KEY))
    documents = []
    for chunk in chunks:
        documents.append({
            "id": chunk["id"], "content": chunk["content"], "filename": chunk["filename"],
            "title": extract_document_title(chunk["filename"]), "chunk_index": chunk["chunk_index"],
            "token_count": chunk["token_count"], "embedding": chunk["embedding"]
        })
    batch_size = 100
    total = len(documents)
    for i in range(0, total, batch_size):
        batch = documents[i:i + batch_size]
        print(f"  Uploading batch {i // batch_size + 1}/{(total + batch_size - 1) // batch_size} ({len(batch)} documents)")
        result = search_client.upload_documents(documents=batch)
        succeeded = sum(1 for r in result if r.succeeded)
        failed = sum(1 for r in result if not r.succeeded)
        print(f"    Succeeded: {succeeded}, Failed: {failed}")

def main():
    parser = argparse.ArgumentParser(description="DisputeAI Document Indexing Pipeline")
    parser.add_argument("--source", choices=["local", "blob"], default="local")
    args = parser.parse_args()
    print("=" * 60)
    print("DisputeAI — Document Indexing Pipeline")
    print("=" * 60)
    print("\n[Step 1/5] Reading documents...")
    if args.source == "local":
        documents = read_local_documents("./documents")
    else:
        documents = read_blob_documents()
    if not documents:
        print("ERROR: No documents found. Exiting.")
        sys.exit(1)
    print("\n[Step 2/5] Chunking documents...")
    all_chunks = []
    for doc in documents:
        chunks = chunk_document(doc)
        print(f"  {doc['filename']}: {len(chunks)} chunks")
        all_chunks.extend(chunks)
    print(f"  Total chunks: {len(all_chunks)}")
    avg_tokens = sum(c["token_count"] for c in all_chunks) / len(all_chunks)
    print(f"  Average tokens per chunk: {avg_tokens:.0f}")
    print("\n[Step 3/5] Generating embeddings...")
    all_chunks = generate_embeddings(all_chunks)
    print(f"  Embeddings generated: {len(all_chunks)}")
    print("\n[Step 4/5] Creating search index...")
    create_search_index()
    print("\n[Step 5/5] Uploading to search index...")
    upload_to_index(all_chunks)
    print("\n" + "=" * 60)
    print("INDEXING COMPLETE")
    print("=" * 60)
    print(f"  Documents processed: {len(documents)}")
    print(f"  Total chunks indexed: {len(all_chunks)}")
    print(f"  Index name: {INDEX_NAME}")
    print(f"  Search endpoint: {SEARCH_ENDPOINT}")
    print("=" * 60)

if __name__ == "__main__":
    main()