add: main.py

This commit is contained in:
2026-06-27 13:52:15 +00:00
commit b5ad3c1028
+146
View File
@@ -0,0 +1,146 @@
import os
import asyncio
import argparse
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
# Load API key from .env
load_dotenv()
# LLM configuration OpenRouter
llm = ChatOpenAI(
model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1",
api_key=os.getenv("OPENAI_API_KEY"),
temperature=0.0,
)
# Pydantic models
class PersonInfo(BaseModel):
name: str = Field(description="Full name of the person")
age: int | None = Field(description="Age of the person, optional", default=None)
profession: str = Field(description="Current profession or role")
skills: list[str] = Field(description="List of skills or technologies")
class MeetingNotes(BaseModel):
date: str = Field(description="Date of the meeting (YYYY-MM-DD)")
participants: list[str] = Field(description="Names of participants")
topics: list[str] = Field(description="Main discussion topics")
decisions: list[str] = Field(description="Decisions made during the meeting")
next_steps: list[str] = Field(description="Action items for next steps")
# Prompt templates
person_prompt = PromptTemplate.from_template(
"""
Extract the following information from the text: name, age (optional), profession, skills.
The text: {input_text}
Output must be in JSON format according to the following schema:
{format_instructions}
"""
)
meeting_prompt = PromptTemplate.from_template(
"""
Extract the following information from the text: date, participants, topics, decisions, next_steps.
The text: {input_text}
Output must be in JSON format according to the following schema:
{format_instructions}
"""
)
# Output parsers
person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)
# Simple heuristic to decide which schema to use
def determine_schema(text: str) -> str:
lower = text.lower()
meeting_keywords = ["meeting", "participants", "date", "topics", "decisions", "next steps", "next_step", "next_step"]
if any(k in lower for k in meeting_keywords):
return "meeting"
return "person"
# Summary generation
def create_summary(parsed: BaseModel, schema: str) -> str:
if schema == "meeting":
return (
f"Meeting on {parsed.date} with participants {', '.join(parsed.participants)}. "
f"Topics: {', '.join(parsed.topics)}. Decisions: {', '.join(parsed.decisions)}. "
f"Next steps: {', '.join(parsed.next_steps)}."
)
else:
age = parsed.age if parsed.age is not None else "N/A"
return (
f"{parsed.name}, age {age}, profession {parsed.profession}. "
f"Skills: {', '.join(parsed.skills)}."
)
# Tool that performs extraction
@tool
def process_text(input_text: str) -> str:
schema = determine_schema(input_text)
if schema == "meeting":
prompt = meeting_prompt
parser = meeting_parser
else:
prompt = person_prompt
parser = person_parser
chain = prompt | llm | parser
parsed = chain.invoke({"input_text": input_text})
summary = create_summary(parsed, schema)
return parsed.model_dump_json(indent=2) + "\n\nSummary: " + summary
# Backend for the agent
backend = CompositeBackend([
LocalShellBackend(workspace_dir="./workspace"),
FilesystemBackend(),
])
# Create the deep agent
agent = create_deep_agent(
model=llm,
tools=[process_text],
backend=backend,
system_prompt="You are a helpful agent that extracts structured data from text.",
)
# CLI entry point
async def main():
parser = argparse.ArgumentParser(description="Extract structured data from text.")
parser.add_argument("--example", choices=["person", "meeting"], help="Run example")
parser.add_argument("--text", type=str, help="Input text")
args = parser.parse_args()
if args.example:
if args.example == "person":
input_text = "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker."
else:
input_text = (
"Дата: 2023-05-01. Участники: Иван, Мария. Темы: проект X, бюджет. "
"Решения: утвердить бюджет. Next steps: подготовить план."
)
elif args.text:
input_text = args.text
else:
input_text = input("Enter text: ")
result = await agent.ainvoke(
{"messages": [HumanMessage(content=input_text)]},
{"configurable": {"thread_id": "session-1"}},
)
output = result["messages"][-1].content
print(output)
if __name__ == "__main__":
asyncio.run(main())