SearchEngine

from agent.rag.search_engine import SearchEngine
from agent.rag.config import Config
from agent.rag.database_manager import DatabaseManager

config = Config()
db_manager = DatabaseManager(config)

search_engine = SearchEngine(config, db_manager)

from agent.rag.search_engine import SearchEngine
from agent.rag.config import Config
from agent.rag.database_manager import DatabaseManager

# Create dependencies
config = Config()
db_manager = DatabaseManager(config)

# Initialize search engine
search_engine = SearchEngine(config, db_manager)

# Basic search
results = search_engine.search_similar_chunks(
    query="What is Apple's revenue?",
    max_results=20
)

# Search specific quarter
results = search_engine.search_similar_chunks(
    query="AI strategy",
    target_quarter="2025_q1"
)

# Search all quarters
results = search_engine.search_similar_chunks(
    query="Financial performance",
    target_quarter="multiple"
)

import asyncio

async def search_with_embedding():
    # Generate embedding asynchronously
    embedding = await search_engine.encode_query_async(
        "What is Microsoft's cloud revenue?"
    )
    
    # Use embedding for custom search
    # embedding.shape -> (1, 384) or (1, 768) depending on model
    return embedding

result = asyncio.run(search_with_embedding())

import asyncio

async def multi_query_search():
    results = await search_engine.search_with_queries_async(
        queries=[
            "What is Apple's revenue growth?",
            "How did Apple perform financially?",
            "Apple revenue trends"
        ],
        target_quarters=['2024_q4', '2025_q1'],
        target_quarter='multiple',
        ticker='AAPL'
    )
    return results

results = asyncio.run(multi_query_search())

import asyncio

async def follow_up():
    results = await search_engine.follow_up_search_async(
        question="What about operating margins?",
        has_tickers=True,
        is_general_question=False,
        is_multi_ticker=False,
        tickers_to_process=['AAPL'],
        target_quarter='2025_q1',
        target_quarters=['2025_q1']
    )
    return results

results = asyncio.run(follow_up())

import asyncio

async def parallel_follow_up():
    results = await search_engine.follow_up_search_parallel_async(
        follow_up_questions=[
            "What about profit margins?",
            "How did expenses change?",
            "What was the guidance?"
        ],
        has_tickers=True,
        is_general_question=False,
        is_multi_ticker=False,
        tickers_to_process=['MSFT'],
        target_quarter='2025_q1',
        target_quarters=['2025_q1']
    )
    return results

results = asyncio.run(parallel_follow_up())

{
    'chunk_text': str,          # The text content of the chunk
    'similarity': float,        # Similarity score (0-1, higher is better)
    'distance': float,          # Distance metric (lower is better)
    'metadata': dict,           # Additional metadata
    'citation': int,            # Chunk index for citation
    'year': int,                # Year of the data
    'quarter': int,             # Quarter number (1-4)
    'ticker': str,              # Stock ticker symbol
    'search_type': str,         # 'vector', 'keyword', or 'hybrid'
    'char_offset': int,         # Character offset in original document
    'chunk_length': int,        # Length of the chunk
    'source_quarter': str       # Source quarter identifier (for multi-quarter)
}

import logging

# Enable search logging
logging.basicConfig(level=logging.INFO)

# Specific loggers:
rag_logger = logging.getLogger('rag_system')
rag_logger.setLevel(logging.INFO)