Multimedika/Bot_Development
0
1import nest_asyncio2import os3from dotenv import load_dotenv4from jinja2 import Template5from pydantic import BaseModel, Field6from pymongo.mongo_client import MongoClient7 8from llama_index.program.openai import OpenAIPydanticProgram9from llama_index.core.extractors import PydanticProgramExtractor10from llama_index.llms.openai import OpenAI11 12from core.prompt import ADD_METADATA_TEMPLATE13from core.summarization.summarizer import SummarizeGenerator14 15nest_asyncio.apply()16 17load_dotenv()18 19 20class NodeMetadata(BaseModel):21 """Metadata for nodes, capturing topic and subtopic from the book."""22 23 topic: str = Field(24 ...,25 description="The main subject or category that the node is associated with, representing a broad theme within the book.",26 )27 subtopic: str = Field(28 ...,29 description="A more specific aspect or section under the main topic, refining the context of the node within the book.",30 )31 32 33def extract_topic(references, content_table):34 uri = os.getenv("MONGO_URI")35 client = MongoClient(uri)36 37 try:38 client.admin.command('ping')39 print("Pinged your deployment. You successfully connected to MongoDB!")40 except Exception as e:41 print(e)42 # Access a specific database43 db = client["summarizer"]44 45 # Access a collection within the database46 collection = db["topic_collection"]47 48 generate_content_table = SummarizeGenerator(references)49 extractor_output, extractor_dics = generate_content_table.extract_content_table(content_table)50 print(extractor_output)51 data_to_insert = {52 "title": references["title"],53 **extractor_dics # Unpack the extractor_output dictionary54 }55 56 collection.insert_one(data_to_insert)57 58 59 add_metadata_template = str(60 Template(ADD_METADATA_TEMPLATE).render(extractor_output=extractor_output)61 )62 63 print("add metadata template : ", add_metadata_template)64 65 llm = OpenAI(temperature=0.1, model="gpt-4o-mini")66 67 openai_program = OpenAIPydanticProgram.from_defaults(68 output_cls=NodeMetadata,69 prompt_template_str="{input}",70 extract_template_str=add_metadata_template,71 llm=llm,72 )73 74 topic_extractor = PydanticProgramExtractor(75 program=openai_program,76 input_key="input",77 show_progress=True,78 extract_template_str=add_metadata_template,79 llm=llm,80 )81 82 return topic_extractor83 