From 3b760b5103975ecdd148ad62f854693e33c468dc Mon Sep 17 00:00:00 2001 From: balabanovan530 <175+balabanovan530@noreply.localhost> Date: Tue, 2 Jun 2026 14:56:37 +0000 Subject: [PATCH] Add agent.py --- agent.py | 173 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 agent.py diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..71daae5 --- /dev/null +++ b/agent.py @@ -0,0 +1,173 @@ +"""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()