From 3638ce76cb3bb7012bb610babfcf3584b464ec02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=92=D0=B8=D0=BA?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BE=D0=B2?= Date: Thu, 2 Jul 2026 05:15:12 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20main.py=20=E2=80=94=20=D0=AD=D0=BA=D0=B7?= =?UTF-8?q?=D0=B0=D0=BC=D0=B5=D0=BD:=20=D0=A1=D1=82=D1=80=D1=83=D0=BA?= =?UTF-8?q?=D1=82=D1=83=D1=80=D0=B8=D1=80=D0=BE=D0=B2=D0=B0=D0=BD=D0=BD?= =?UTF-8?q?=D1=8B=D0=B9=20=D0=B2=D1=8B=D0=B2=D0=BE=D0=B4=20(Pydantic)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 277 ++++++++++++++++++++++++++++---------------------------- 1 file changed, 140 insertions(+), 137 deletions(-) diff --git a/main.py b/main.py index ad469f7..225b023 100644 --- a/main.py +++ b/main.py @@ -1,19 +1,17 @@ import os import asyncio -import sys from typing import Union from dotenv import load_dotenv from pydantic import BaseModel, Field -from langchain_core.prompts import PromptTemplate -from langchain_core.output_parsers import PydanticOutputParser from langchain_openai import ChatOpenAI from langchain_core.messages import HumanMessage +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 from deepagents.tools import tool -# Load environment variables (API key) load_dotenv() # ---------------------------------------------------------------------- @@ -22,60 +20,22 @@ load_dotenv() class PersonInfo(BaseModel): name: str = Field(description="Full name of the person") age: Union[int, None] = Field( - default=None, - description="Age of the person, if mentioned" + default=None, description="Age in years, optional if not mentioned" ) profession: str = Field(description="Professional title or occupation") - skills: list[str] = Field( - description="List of skills or technologies mentioned" - ) + skills: list[str] = Field(description="List of key skills or technologies") class MeetingNotes(BaseModel): - date: str = Field(description="Date of the meeting in ISO format or natural language") - participants: list[str] = Field(description="Names of participants") - topics: list[str] = Field(description="Main topics discussed") - decisions: list[str] = Field(description="Key decisions made") + date: str = Field(description="Date of the meeting in ISO format (YYYY-MM-DD)") + participants: list[str] = Field(description="List of participant names") + 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 or next steps") # ---------------------------------------------------------------------- -# Output parsers -# ---------------------------------------------------------------------- -person_parser = PydanticOutputParser(pydantic_object=PersonInfo) -meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes) - -# ---------------------------------------------------------------------- -# Prompt templates -# ---------------------------------------------------------------------- -person_prompt = PromptTemplate.from_template( - """Extract the following information about a person from the given text. -Return the data in JSON format that matches the provided schema. - -{format_instructions} - -Text: -\"\"\" -{input_text} -\"\"\" -""" -) - -meeting_prompt = PromptTemplate.from_template( - """Extract structured meeting notes from the given text. -Return the data in JSON format that matches the provided schema. - -{format_instructions} - -Text: -\"\"\" -{input_text} -\"\"\" -""" -) - -# ---------------------------------------------------------------------- -# LLM configuration (OpenRouter) +# LLM configuration (OpenRouter via langchain-openai) # ---------------------------------------------------------------------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", @@ -85,113 +45,156 @@ llm = ChatOpenAI( ) # ---------------------------------------------------------------------- -# DeepAgents setup (required by the course) +# Prompt templates with format instructions # ---------------------------------------------------------------------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) +person_parser = PydanticOutputParser(pydantic_object=PersonInfo) +meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes) + +person_prompt = PromptTemplate.from_template( + """Extract the following information about a person and output it as JSON that matches the given schema. + +{format_instructions} + +Text: +{input_text} +""" +) + +meeting_prompt = PromptTemplate.from_template( + """Extract structured meeting notes from the given text and output them as JSON that matches the given schema. + +{format_instructions} + +Text: +{input_text} +""" +) + +# ---------------------------------------------------------------------- +# Simple routing based on keyword heuristics +# ---------------------------------------------------------------------- +def select_schema(text: str) -> str: + """Return 'person' or 'meeting' depending on the content.""" + lowered = text.lower() + meeting_keywords = ["встреча", "meeting", "участники", "participants", "agenda", "решения"] + if any(word in lowered for word in meeting_keywords): + return "meeting" + return "person" + + +# ---------------------------------------------------------------------- +# DeepAgent creation (required by the course) +# ---------------------------------------------------------------------- +backend = CompositeBackend( + [ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), + ] +) @tool -def dummy_tool(query: str) -> str: - """Placeholder tool required by the agent; simply echoes the query.""" - return f"Echo: {query}" +def route_and_process(text: str) -> str: + """Detect the type of the input text, run the appropriate extraction chain, + and return the JSON representation of the validated Pydantic model.""" + schema_type = select_schema(text) + + if schema_type == "person": + parser = person_parser + prompt = person_prompt + else: + parser = meeting_parser + prompt = meeting_prompt + + chain = ( + prompt.partial(format_instructions=parser.get_format_instructions()) + | llm + | parser + ) + result = chain.invoke({"input_text": text}) + # Return pretty JSON for CLI display + return result.model_dump_json(indent=2) + agent = create_deep_agent( model=llm, - tools=[dummy_tool], + tools=[route_and_process], backend=backend, system_prompt="You are a helpful assistant that extracts structured data.", ) -# ---------------------------------------------------------------------- -# Routing logic -# ---------------------------------------------------------------------- -def choose_schema(text: str) -> str: - """ - Simple heuristic to decide which schema to use. - If the text contains keywords typical for meeting notes, use MeetingNotes, - otherwise assume it describes a person. - """ - meeting_keywords = ["встреча", "meeting", "участники", "participants", "agenda", "решения", "decisions"] - lowered = text.lower() - for kw in meeting_keywords: - if kw in lowered: - return "meeting" - return "person" - -# DESIGN DECISION: Use a keyword-based heuristic for schema selection. -# NECESSITY: The assignment requires a routing step but does not mandate a sophisticated classifier. -# OPTIMALITY: This approach is fast, deterministic, and does not require additional model calls. -# ALTERNATIVES CONSIDERED: A separate classification LLM call was considered but would increase latency -# and cost without adding educational value for this simple task. # ---------------------------------------------------------------------- -# Extraction functions +# CLI interface # ---------------------------------------------------------------------- -async def extract_person(text: str) -> PersonInfo: - prompt = person_prompt.partial_variables({ - "format_instructions": person_parser.get_format_instructions(), - "input_text": text, - }) - response = await llm.ainvoke([HumanMessage(content=prompt.format())]) - parsed = person_parser.parse(response.content) - return parsed +EXAMPLES = { + "person": "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker.", + "meeting": """Дата: 2024-09-15 +Участники: Иван, Мария, Алексей +Темы: План проекта, бюджет, сроки +Решения: Утвердить бюджет в 500k, начать разработку 1 октября +Следующие шаги: Иван подготовит ТЗ, Мария соберёт требования, Алексей настроит окружение.""" +} -async def extract_meeting(text: str) -> MeetingNotes: - prompt = meeting_prompt.partial_variables({ - "format_instructions": meeting_parser.get_format_instructions(), - "input_text": text, - }) - response = await llm.ainvoke([HumanMessage(content=prompt.format())]) - parsed = meeting_parser.parse(response.content) - return parsed - -# ---------------------------------------------------------------------- -# Main CLI -# ---------------------------------------------------------------------- -async def main(): - if len(sys.argv) > 1: - input_text = " ".join(sys.argv[1:]) - else: - print("Enter the text (finish with an empty line):") - lines = [] - while True: - line = input() - if line == "": - break - lines.append(line) - input_text = "\n".join(lines) - - schema = choose_schema(input_text) - - # Use the deep agent as a required component; we invoke it with a trivial message. - # The result is not used for extraction but satisfies the course requirement. - await agent.ainvoke( - {"messages": [HumanMessage(content="Prepare for extraction")]}, - {"configurable": {"thread_id": "session-cli"}}, +async def run_example(example_key: str): + text = EXAMPLES[example_key] + result = await agent.ainvoke( + {"messages": [HumanMessage(content=text)]}, + {"configurable": {"thread_id": f"example-{example_key}"}}, ) + print("Input text:") + print(text) + print("\nExtracted JSON:") + print(result["messages"][-1].content) - if schema == "person": - result = await extract_person(input_text) - else: - result = await extract_meeting(input_text) - # Output the raw model_dump and a short summary - print("\n--- Structured Output (model_dump) ---") - print(result.model_dump()) - print("\n--- Summary ---") - if isinstance(result, PersonInfo): - summary = f"{result.name}, {result.age or 'N/A'} years old, works as {result.profession}. Skills: {', '.join(result.skills)}." - else: - summary = ( - f"Meeting on {result.date} with {', '.join(result.participants)}. " - f"Topics: {', '.join(result.topics)}. Decisions: {', '.join(result.decisions)}. " - f"Next steps: {', '.join(result.next_steps)}." +def interactive_mode(): + print("Enter text (empty line to finish):") + lines = [] + while True: + line = input() + if line == "": + break + lines.append(line) + user_text = "\n".join(lines) + if not user_text.strip(): + print("No input provided.") + return + + result = asyncio.run( + agent.ainvoke( + {"messages": [HumanMessage(content=user_text)]}, + {"configurable": {"thread_id": "interactive-session"}}, ) - print(summary) + ) + print("\nExtracted JSON:") + print(result["messages"][-1].content) + + +def main(): + import argparse + + parser = argparse.ArgumentParser(description="Structured extraction demo") + group = parser.add_mutually_exclusive_group() + group.add_argument( + "--example", + choices=["person", "meeting"], + help="Run a built-in example", + ) + group.add_argument( + "--interactive", + action="store_true", + help="Enter interactive mode", + ) + args = parser.parse_args() + + if args.example: + asyncio.run(run_example(args.example)) + elif args.interactive: + interactive_mode() + else: + parser.print_help() if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + main() \ No newline at end of file