add: main.py
This commit is contained in:
@@ -0,0 +1,112 @@
|
||||
import os
|
||||
import asyncio
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
from langchain_core.output_parsers import PydanticOutputParser
|
||||
from pydantic import BaseModel, Field
|
||||
from deepagents import create_deep_agent
|
||||
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
|
||||
|
||||
# ---------- LLM ----------
|
||||
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(default=None, description="Age in years, optional")
|
||||
profession: str = Field(description="Current profession or job title")
|
||||
skills: list[str] = Field(description="List of professional skills")
|
||||
|
||||
class MeetingNotes(BaseModel):
|
||||
date: str = Field(description="Meeting date in ISO format (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(
|
||||
input_variables=["text"],
|
||||
template="""
|
||||
You are a data extraction assistant. Extract the following information from the given text and output a JSON object that matches the PersonInfo schema.
|
||||
|
||||
Text: {text}
|
||||
|
||||
Output must be a valid JSON object with fields: name, age, profession, skills.
|
||||
"""
|
||||
)
|
||||
|
||||
meeting_prompt = PromptTemplate(
|
||||
input_variables=["text"],
|
||||
template="""
|
||||
You are a data extraction assistant. Extract the following information from the given text and output a JSON object that matches the MeetingNotes schema.
|
||||
|
||||
Text: {text}
|
||||
|
||||
Output must be a valid JSON object with fields: date, participants, topics, decisions, next_steps.
|
||||
"""
|
||||
)
|
||||
|
||||
# ---------- Parsers ----------
|
||||
person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
|
||||
meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)
|
||||
|
||||
# ---------- Chains ----------
|
||||
person_chain = person_prompt | llm | person_parser
|
||||
meeting_chain = meeting_prompt | llm | meeting_parser
|
||||
|
||||
# ---------- Backend ----------
|
||||
backend = CompositeBackend([
|
||||
LocalShellBackend(workspace_dir="./workspace"),
|
||||
FilesystemBackend(),
|
||||
])
|
||||
|
||||
# ---------- Agent ----------
|
||||
agent = create_deep_agent(
|
||||
model=llm,
|
||||
tools=[],
|
||||
backend=backend,
|
||||
system_prompt="You are a structured data extraction agent. Decide whether the input text is about a person or a meeting and return the parsed JSON accordingly.",
|
||||
)
|
||||
|
||||
# ---------- Helper for routing ----------
|
||||
async def route_and_parse(text: str):
|
||||
# Simple heuristic: if the word "meeting" or "встреча" appears, treat as meeting
|
||||
if "meeting" in text.lower() or "встреча" in text.lower():
|
||||
result = await meeting_chain.ainvoke({"text": text})
|
||||
return MeetingNotes(**result)
|
||||
else:
|
||||
result = await person_chain.ainvoke({"text": text})
|
||||
return PersonInfo(**result)
|
||||
|
||||
# ---------- CLI ----------
|
||||
async def main():
|
||||
examples = {
|
||||
"person": "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker.",
|
||||
"meeting": "Встреча 2026-06-10. Участники: Иван, Мария. Темы: проект X, бюджет. Решения: утвердить план. Next steps: подготовить презентацию.",
|
||||
}
|
||||
print("Выберите пример: 1 - человек, 2 - встреча, 3 - ввод с клавиатуры")
|
||||
choice = input("> ")
|
||||
if choice == "1":
|
||||
text = examples["person"]
|
||||
elif choice == "2":
|
||||
text = examples["meeting"]
|
||||
else:
|
||||
text = input("Введите текст: ")
|
||||
parsed = await route_and_parse(text)
|
||||
print("\nРезультат (model_dump):")
|
||||
print(parsed.model_dump(indent=2))
|
||||
print("\nКраткая сводка:")
|
||||
if isinstance(parsed, PersonInfo):
|
||||
print(f"{parsed.name}, {parsed.age or 'неизвестно'} лет, {parsed.profession}. Навыки: {', '.join(parsed.skills)}")
|
||||
else:
|
||||
print(f"Дата: {parsed.date}\nУчастники: {', '.join(parsed.participants)}\nТемы: {', '.join(parsed.topics)}\nРешения: {', '.join(parsed.decisions)}\nСледующие шаги: {', '.join(parsed.next_steps)}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user