Add agent.py

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()