diff --git a/main.py b/main.py new file mode 100644 index 0000000..8cc22a5 --- /dev/null +++ b/main.py @@ -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() \ No newline at end of file