RAGAgent

​

Overview

​

Class Definition

​

RAGAgent

from agent.rag.rag_agent import RAGAgent

agent = RAGAgent(openai_api_key="your-api-key")

​

Constructor

​

__init__(openai_api_key=None)

• ValueError - If OpenAI API key is not provided and not found in environment

import os
from agent.rag.rag_agent import RAGAgent

# With explicit API key
agent = RAGAgent(openai_api_key="sk-...")

# From environment variable
agent = RAGAgent()  # Uses OPENAI_API_KEY env var

​

Core Methods

​

set_database_connection(db_connection)

from psycopg2 import connect

# Create database connection
conn = connect(database="finance_db", user="user")

# Set on agent
agent.set_database_connection(conn)

​

RAG Flow Pipeline

​

Stage 1: Setup

• Initialize context and logging
• Set maximum iterations

​

Stage 2: Reasoning

• Analyze question with LLM
• Extract metadata (tickers, quarters)
• Explain research approach

​

Stage 3: Search Planning

• Resolve temporal references to specific quarters
• Generate queries for each data source

​

Stage 4: Parallel Multi-Source Search

# Automatically triggered if news is relevant
# Returns recent news articles and context

# Searches annual reports and filings
# Extracts financial data and disclosures

# Hybrid search combining:
# - Vector similarity search
# - Keyword-based search
# - Merged and ranked results

​

Stage 5: Context Preparation

• Build news/10-K context strings
• Combine citations from all sources

​

Stage 6: Improvement Loop

• Generate initial answer
• Evaluate quality and confidence
• Perform follow-up searches if needed
• Re-generate answer with additional context

​

Stage 7: Finalization

• Deduplicate citations
• Build final result structure
• Update conversation memory

​

Internal Methods

​

Question Analysis

async def _perform_question_analysis(question: str, conversation_id: str) -> tuple

• question_analysis (Dict): Analysis results including status, tickers, question type, processed question, and quarter context
• target_quarters (List[str]): Quarters to search (e.g., ['2025_q1', '2025_q2'])
• error_dict (Dict): Error information if validation fails

​

Search Execution

async def _execute_search(ctx: RAGFlowContext) -> None

• question
• question_analysis
• target_quarters

• individual_results
• all_chunks
• all_citations
• search_time
• is_general_question
• is_multi_ticker
• tickers_to_process
• target_quarter

​

Ticker Processing

def _process_single_ticker_sync(
    ticker: str,
    question: str,
    processed_question: str,
    is_multi_ticker: bool,
    target_quarters: List[str]
) -> Dict[str, Any]

​

Parallel Processing

def _process_tickers_parallel_sync(
    tickers: List[str],
    question: str,
    processed_question: str,
    is_multi_ticker: bool,
    target_quarters: List[str]
) -> List[Dict[str, Any]]

​

Configuration

​

Usage Examples

from agent.rag.rag_agent import RAGAgent
import os

# Initialize agent
agent = RAGAgent(openai_api_key=os.getenv("OPENAI_API_KEY"))

# The agent is ready to process questions
# Question analysis and search will happen automatically
# when the agent's execute method is called

from agent.rag.rag_agent import RAGAgent
from psycopg2 import connect
import os

# Initialize agent
agent = RAGAgent(openai_api_key=os.getenv("OPENAI_API_KEY"))

# Set up database connection for conversation memory
conn = connect(
    database="finance_db",
    user="user",
    password="password",
    host="localhost"
)

agent.set_database_connection(conn)

# Agent can now access conversation history

from agent.rag.rag_agent import RAGAgent
from agent.rag.config import Config
import os

# Create custom config
config = Config()
config.set("max_quarters", 8)  # Limit to 2 years
config.set("max_tickers", 6)   # Allow more tickers
config.set("vector_weight", 0.8)  # Prefer vector search

# Initialize agent
agent = RAGAgent(openai_api_key=os.getenv("OPENAI_API_KEY"))
agent.config = config

​

Components

​

Logging

import logging

# Configure logging to see agent operations
logging.basicConfig(level=logging.INFO)

# Specific loggers:
# - 'rag_agent' - RAG flow operations
# - 'rag_system' - Low-level system operations

​

Error Handling

• Missing API Key: Raises ValueError on initialization
• Database Errors: Logs warning and continues with default config
• Quarter Unavailable: Returns user-friendly error with available quarters
• No Results: Returns empty context response when no data is found
• Search Failures: Falls back gracefully and logs errors

​

Performance Optimization

1. Parallel Search: Multiple data sources searched concurrently
2. Thread Pool: CPU-bound operations executed in thread pool
3. Multi-Ticker Parallel: Tickers processed simultaneously
4. Multi-Quarter Parallel: Quarters searched in parallel
5. Async Operations: Non-blocking I/O operations

​

Related