setup_documents.py

"""
Setup script to populate LanceDB with sample documents for the Voice Bot
"""
import os
from lancedb_service import lancedb_service
import asyncio
import logging

logger = logging.getLogger("voicebot")

# Sample government documents content
SAMPLE_DOCUMENTS = [
    {
        "content": """Pension is a regular payment made during a person's retirement from an investment fund to which that person or their employer has contributed during their working career. In the context of government employees in India, pension refers to the retirement benefits provided to civil servants after completing their service period.

The pension system for government employees in India includes:
1. Basic Pension: Calculated based on the last drawn salary and years of service
2. Dearness Relief (DR): Additional amount to counter inflation
3. Medical Benefits: Healthcare coverage post-retirement
4. Family Pension: Benefits provided to family members after the employee's death

Key features:
- Minimum service period: Usually 10 years for qualifying for pension
- Calculation formula: (Last drawn basic salary × years of service) / 70 (for pre-2016 employees)
- Monthly payment: Pension is paid monthly to the retired employee
- Indexation: Pension amount is revised periodically to account for inflation""",
        "filename": "pension_basics.txt",
        "source": "Government Pension Manual"
    },
    {
        "content": """Dearness Allowance (DA) is a component of salary paid to government employees and pensioners to offset the impact of inflation. For government employees in Rajasthan and other states, DA is calculated as a percentage of basic salary.

Current DA rates and impact:
- DA is revised twice a year (January and July)
- Based on All India Consumer Price Index (AICPI)
- Recent DA increment: 6% increase effective from January 2024
- Impact on salary: For an employee with basic salary of ₹50,000, a 6% DA increment adds ₹3,000 per month
- Impact on pension: Pensioners also receive corresponding Dearness Relief (DR) increase

Calculation method:
- DA percentage = ((Average AICPI for 12 months - Base AICPI) / Base AICPI) × 100
- Rounded to nearest 0.5%

Benefits of 6% DA increment:
1. Increased monthly income for serving employees
2. Higher pension amounts for retirees
3. Improved purchasing power to counter inflation
4. Applicable to all government employees across pay scales""",
        "filename": "dearness_allowance_impact.txt", 
        "source": "DA Revision Circular 2024"
    },
    {
        "content": """Government employees in Rajasthan are entitled to various retirement benefits including pension, gratuity, and provident fund. The pension system has undergone reforms with the introduction of the New Pension Scheme (NPS) for employees joining after 2004.

Rajasthan Pension Rules:
1. Old Pension Scheme (OPS): For employees joined before 2004
   - Defined benefit scheme
   - Pension = 50% of last drawn salary after 33 years of service
   - Family pension available to spouse and children

2. New Pension Scheme (NPS): For employees joined after 2004
   - Defined contribution scheme
   - Employee contribution: 10% of basic salary
   - Government contribution: 14% of basic salary
   - Market-linked returns

Recent Developments:
- Rajasthan government restored OPS for all employees in 2022
- Employees under NPS can now opt for OPS benefits
- Enhanced pension benefits for teachers and other government employees

Pension Processing:
- Application to be submitted 6 months before retirement
- Required documents: Service records, medical fitness certificate, family details
- Processing time: Usually 3-6 months
- Monthly pension credit: Through NEFT to bank account""",
        "filename": "rajasthan_pension_rules.txt",
        "source": "Rajasthan Government Pension Manual 2024"
    }
]

async def setup_sample_documents():
    """Create sample documents in the database"""
    try:
        logger.info("🗂️ Setting up sample government documents...")
        
        # Create document objects that match LanceDB expectations
        docs = []
        for doc_data in SAMPLE_DOCUMENTS:
            # Create a simple document object with the expected attributes
            doc = type('Document', (), {
                'page_content': doc_data["content"],
                'metadata': {"source": doc_data["source"], "filename": doc_data["filename"]}
            })()
            docs.append(doc)
        
        # Add to government_docs knowledge base
        await lancedb_service.add_documents(
            docs=docs,
            user_id="system",
            knowledge_base="government_docs", 
            filename="sample_documents.txt"
        )
        
        # Also create a specific rajasthan_documents table entry
        # Check if rajasthan_documents table exists, create if not
        try:
            if "rajasthan_documents" not in lancedb_service.db.table_names():
                # Create the table using the same structure as documents table
                import pandas as pd
                import uuid
                from datetime import datetime
                
                sample_data = pd.DataFrame({
                    "id": [str(uuid.uuid4())],
                    "content": ["sample"],
                    "filename": ["sample"],
                    "vector": [lancedb_service.embedding_model.embed_query("sample")]
                })
                lancedb_service.db.create_table("rajasthan_documents", sample_data)
                # Delete sample data
                tbl = lancedb_service.db.open_table("rajasthan_documents")
                tbl.delete("id = 'sample'")
                
                # Now add the real documents
                rajasthan_docs = []
                for doc_data in SAMPLE_DOCUMENTS:
                    embedding = lancedb_service.embedding_model.embed_query(doc_data["content"])
                    raj_doc = {
                        "id": str(uuid.uuid4()),
                        "content": doc_data["content"],
                        "filename": doc_data["filename"],
                        "vector": embedding
                    }
                    rajasthan_docs.append(raj_doc)
                
                df = pd.DataFrame(rajasthan_docs)
                tbl.add(df)
                
        except Exception as e:
            logger.warning(f"⚠️ Could not setup rajasthan_documents table: {e}")
            
        logger.info("✅ Sample documents added successfully")
        
    except Exception as e:
        logger.error(f"❌ Error setting up documents: {e}")

if __name__ == "__main__":
    asyncio.run(setup_sample_documents())