68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
import os
|
|
import sys
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.prompts import PromptTemplate
|
|
from langchain_core.output_parsers import PydanticOutputParser
|
|
from models import PersonInfo, MeetingNotes
|
|
|
|
# Choose LLM
|
|
llm = ChatOpenAI(temperature=0)
|
|
|
|
# Prepare parser for each model
|
|
person_parser = PydanticOutputParser(pydantic_object=PersonInfo)
|
|
meeting_parser = PydanticOutputParser(pydantic_object=MeetingNotes)
|
|
|
|
# Prompt templates
|
|
person_prompt = PromptTemplate(
|
|
input_variables=["text"],
|
|
template="""
|
|
Extract structured data about a person from the following text.
|
|
The output must follow the JSON schema:
|
|
{schema}
|
|
Text: {text}
|
|
""",
|
|
partial_variables={"schema": person_parser.get_format_instructions()},
|
|
)
|
|
|
|
meeting_prompt = PromptTemplate(
|
|
input_variables=["text"],
|
|
template="""
|
|
Extract structured data about a meeting from the following text.
|
|
The output must follow the JSON schema:
|
|
{schema}
|
|
Text: {text}
|
|
""",
|
|
partial_variables={"schema": meeting_parser.get_format_instructions()},
|
|
)
|
|
|
|
# Helper to detect type
|
|
|
|
def detect_type(text: str) -> str:
|
|
# Simple heuristic: if contains "meeting" or "participants" -> meeting, else person
|
|
lower = text.lower()
|
|
if "meeting" in lower or "participants" in lower or "topics" in lower:
|
|
return "meeting"
|
|
return "person"
|
|
|
|
def main():
|
|
if len(sys.argv) > 1:
|
|
input_text = " ".join(sys.argv[1:])
|
|
else:
|
|
print("Enter text (or press Ctrl-D to exit):")
|
|
input_text = sys.stdin.read().strip()
|
|
if not input_text:
|
|
print("No input provided.")
|
|
return
|
|
|
|
t = detect_type(input_text)
|
|
if t == "person":
|
|
chain = person_prompt | llm | person_parser
|
|
result = chain.invoke({"text": input_text})
|
|
print("\nParsed PersonInfo:\n", result.model_dump(indent=2))
|
|
else:
|
|
chain = meeting_prompt | llm | meeting_parser
|
|
result = chain.invoke({"text": input_text})
|
|
print("\nParsed MeetingNotes:\n", result.model_dump(indent=2))
|
|
|
|
if __name__ == "__main__":
|
|
main() |