# in your agent definition file, e.g. rag/doc_parse_agent.py

import os, base64, json
from typing import Dict, Any
from langchain.tools import StructuredTool
from langchain.agents import initialize_agent, AgentType
from langchain_openai import ChatOpenAI
from rag.ocr_mistral import parse_passport
from rag.ocr_azure import parse_passport_azure

# -- LangChain tool wrapper -------------------------------------------------
azure_passport_tool = StructuredTool.from_defaults(
    name="parse_passport_azure",
    func=parse_passport_azure,
    description="Extract key fields from raw passport bytes.",
)

# -- Agent factory ----------------------------------------------------------
def create_passport_agent():
    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
    return initialize_agent(
        tools=[azure_passport_tool],
        llm=llm,
        agent=AgentType.STRUCTURED_TOOLING,
        verbose=False,
        handle_parsing_errors=True,
    )

# quick demo when executed directly
if __name__ == "__main__":
    ag = create_passport_agent()
    with open("data/sample_passport.png", "rb") as f:
        result = ag.run(
            "Extract passport info",
            input={"parse_passport_azure": f.read()}
        )
    print("Agent returned:", result)