# gen-rag-crawl/agent_utils.py
from pydantic_ai_agent import pydantic_ai_agent, PydanticAIDeps
from db import get_chroma_client, init_collection
from openai import AsyncOpenAI
import os
from dotenv import load_dotenv

load_dotenv()
openai_client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))

chroma_collection = init_collection()

async def run_agent_with_streaming(user_input: str):
    """Run the agent with streaming text for the user_input prompt."""
    deps = PydanticAIDeps(collection=chroma_collection, openai_client=openai_client)

    async with pydantic_ai_agent.run_stream(
        user_input,
        deps=deps,
        message_history=[]  # Remove Streamlit dependency
    ) as result:
        partial_text = ""
        async for chunk in result.stream_text(delta=True):
            partial_text += chunk
        return partial_text  # Return plain text instead of Streamlit objects