Add main.py
This commit is contained in:
@@ -0,0 +1,68 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user