task_6a1865008a94f887e50d471c_auto/agent.py

"""CLI tool that extracts structured data from text using LangChain and Pydantic.

The tool supports two schemas:
- PersonInfo
- MeetingNotes

It automatically detects which schema to use based on the input text.
"""

import sys
import argparse
from typing import List

from dotenv import load_dotenv
import os

from pydantic import BaseModel, Field

from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser

# Load environment variables
load_dotenv()

# LLM configuration – user can override via env vars
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-3.5-turbo")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "")

# ---------- Pydantic models ----------
class PersonInfo(BaseModel):
    """Information about a person."""
    name: str = Field(..., description="Person's name")
    age: int = Field(..., description="Age of the person")
    profession: str = Field(..., description="Profession of the person")
    skills: List[str] = Field(..., description="List of skills")

class MeetingNotes(BaseModel):
    """Notes from a meeting."""
    title: str = Field(..., description="Title of the meeting")
    date: str = Field(..., description="Date of the meeting")
    participants: List[str] = Field(..., description="List of participants")
    agenda: List[str] = Field(..., description="Agenda items")

# ---------- LLM instance ----------
llm = ChatOpenAI(
    model=OPENAI_MODEL,
    api_key=OPENAI_API_KEY,
    base_url=OPENAI_BASE_URL if OPENAI_BASE_URL else None,
    temperature=0,
)

# ---------- Prompts and parsers ----------
# 1. Schema detection prompt
schema_detection_prompt = PromptTemplate(
    input_variables=["text"],
    template="""
Determine whether the following text is about a person or a meeting. Respond with one of the words:
- PERSON
- MEETING

Text: {text}
Answer:
"""
)

# 2. PersonInfo extraction prompt
person_prompt = PromptTemplate(
    input_variables=["text"],
    template="""
Extract the following information about a person from the text:
- name
- age
- profession
- skills (comma separated list)

Return a JSON object with fields name, age, profession, skills.

Text: {text}
JSON:
"""
)
person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
person_chain = person_prompt | llm | person_parser

# 3. MeetingNotes extraction prompt
meeting_prompt = PromptTemplate(
    input_variables=["text"],
    template="""
Extract the following information about a meeting from the text:
- title
- date
- participants (comma separated list)
- agenda (comma separated list)

Return a JSON object with fields title, date, participants, agenda.

Text: {text}
JSON:
"""
)
meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)
meeting_chain = meeting_prompt | llm | meeting_parser

# ---------- Helper functions ----------

def detect_schema(text: str) -> str:
    """Return "PERSON" or "MEETING" based on LLM classification."""
    try:
        result = schema_detection_prompt.invoke({"text": text}, llm=llm)
        # The result is a string; strip and uppercase
        return result.strip().upper()
    except Exception as e:
        raise RuntimeError(f"Schema detection failed: {e}")


def extract_person(text: str) -> PersonInfo:
    """Run the person extraction chain."""
    try:
        return person_chain.invoke({"text": text})
    except Exception as e:
        raise RuntimeError(f"Person extraction failed: {e}")


def extract_meeting(text: str) -> MeetingNotes:
    """Run the meeting extraction chain."""
    try:
        return meeting_chain.invoke({"text": text})
    except Exception as e:
        raise RuntimeError(f"Meeting extraction failed: {e}")

# ---------- CLI ----------

def main() -> None:
    parser = argparse.ArgumentParser(description="Extract structured data from text.")
    parser.add_argument(
        "text",
        nargs="?",
        help="Input text. If omitted, read from stdin.",
    )
    args = parser.parse_args()

    if args.text:
        input_text = args.text
    else:
        input_text = sys.stdin.read().strip()
        if not input_text:
            print("No input provided.", file=sys.stderr)
            sys.exit(1)

    # Detect schema
    schema_type = detect_schema(input_text)
    if schema_type == "PERSON":
        obj = extract_person(input_text)
    elif schema_type == "MEETING":
        obj = extract_meeting(input_text)
    else:
        print(f"Unable to determine schema type: {schema_type}", file=sys.stderr)
        sys.exit(1)

    # Output
    print(obj.model_dump(indent=2))
    # Simple summary
    if isinstance(obj, PersonInfo):
        summary = f"{obj.name}, {obj.age} years old, works as {obj.profession}.\n" + f"Skills: {', '.join(obj.skills)}."
    else:
        summary = f"Meeting '{obj.title}' on {obj.date} with participants {', '.join(obj.participants)}. Agenda: {', '.join(obj.agenda)}."
    print("\nSummary:\n" + summary)


if __name__ == "__main__":
    main()