147 lines
4.9 KiB
Python
147 lines
4.9 KiB
Python
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())
|