fix: main.py — Экзамен: Структурированный вывод (Pydantic)

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,
+        default=None, description="Age in years, optional if not mentioned"
-        description="Age of the person, if mentioned"
     )
     profession: str = Field(description="Professional title or occupation")
-    skills: list[str] = Field(
+    skills: list[str] = Field(description="List of key skills or technologies")
-        description="List of skills or technologies mentioned"
-    )
 
 
 class MeetingNotes(BaseModel):
-    date: str = Field(description="Date of the meeting in ISO format or natural language")
+    date: str = Field(description="Date of the meeting in ISO format (YYYY-MM-DD)")
-    participants: list[str] = Field(description="Names of participants")
+    participants: list[str] = Field(description="List of participant names")
-    topics: list[str] = Field(description="Main topics discussed")
+    topics: list[str] = Field(description="Main discussion topics")
-    decisions: list[str] = Field(description="Key decisions made")
+    decisions: list[str] = Field(description="Decisions made during the meeting")
     next_steps: list[str] = Field(description="Action items or next steps")
 
 
 # ----------------------------------------------------------------------
-# Output parsers
+# LLM configuration (OpenRouter via langchain-openai)
-# ----------------------------------------------------------------------
-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 = 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([
+person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
-    LocalShellBackend(workspace_dir="./workspace"),
+meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)
-    FilesystemBackend(),
+
-])
+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:
+def route_and_process(text: str) -> str:
-    """Placeholder tool required by the agent; simply echoes the query."""
+    """Detect the type of the input text, run the appropriate extraction chain,
-    return f"Echo: {query}"
+    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:
+EXAMPLES = {
-    prompt = person_prompt.partial_variables({
+    "person": "Анна, 28 лет, Python-разработчик. Навыки: FastAPI, Docker.",
-        "format_instructions": person_parser.get_format_instructions(),
+    "meeting": """Дата: 2024-09-15
-        "input_text": text,
+Участники: Иван, Мария, Алексей
-    })
+Темы: План проекта, бюджет, сроки
-    response = await llm.ainvoke([HumanMessage(content=prompt.format())])
+Решения: Утвердить бюджет в 500k, начать разработку 1 октября
-    parsed = person_parser.parse(response.content)
+Следующие шаги: Иван подготовит ТЗ, Мария соберёт требования, Алексей настроит окружение."""
-    return parsed
+}
 
 
-async def extract_meeting(text: str) -> MeetingNotes:
+async def run_example(example_key: str):
-    prompt = meeting_prompt.partial_variables({
+    text = EXAMPLES[example_key]
-        "format_instructions": meeting_parser.get_format_instructions(),
+    result = await agent.ainvoke(
-        "input_text": text,
+        {"messages": [HumanMessage(content=text)]},
-    })
+        {"configurable": {"thread_id": f"example-{example_key}"}},
-    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"}},
     )
+    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
+def interactive_mode():
-    print("\n--- Structured Output (model_dump) ---")
+    print("Enter text (empty line to finish):")
-    print(result.model_dump())
+    lines = []
-    print("\n--- Summary ---")
+    while True:
-    if isinstance(result, PersonInfo):
+        line = input()
-        summary = f"{result.name}, {result.age or 'N/A'} years old, works as {result.profession}. Skills: {', '.join(result.skills)}."
+        if line == "":
-    else:
+            break
-        summary = (
+        lines.append(line)
-            f"Meeting on {result.date} with {', '.join(result.participants)}. "
+    user_text = "\n".join(lines)
-            f"Topics: {', '.join(result.topics)}. Decisions: {', '.join(result.decisions)}. "
+    if not user_text.strip():
-            f"Next steps: {', '.join(result.next_steps)}."
+        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())
+    main()