ASE Course Text Article

text_article.py

import chromadb
from llama_index.embeddings.openai import OpenAIEmbedding
import openai
from llama_index.llms.openai import OpenAI as LlamaOpenAI
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core.indices.vector_store import VectorStoreIndex
from llama_index.core.vector_stores import FilterOperator, FilterCondition, MetadataFilter, MetadataFilters
from llama_index.core.query_engine import CustomQueryEngine
from llama_index.core.retrievers import BaseRetriever
from llama_index.core import get_response_synthesizer
from llama_index.core.response_synthesizers import BaseSynthesizer
from llama_index.core import PromptTemplate
from dotenv import load_dotenv
from os import getenv

load_dotenv()

OPENAI_API_KEY = getenv("OPENAI_API_KEY")
CHROMA_DB_HOST = "0.0.0.0"

embed_model = OpenAIEmbedding(model="text-embedding-3-small", api_key=OPENAI_API_KEY)
openai_llm = LlamaOpenAI(model="o1-mini", api_key=OPENAI_API_KEY)

db = chromadb.HttpClient(CHROMA_DB_HOST)

collection = db.get_collection(name="course_content")
vector_store = ChromaVectorStore(chroma_collection = collection)
index = VectorStoreIndex.from_vector_store(vector_store)
synthesizer = get_response_synthesizer(response_mode="refine")
retriever = index.as_retriever(filters=MetadataFilters(filters=[
                    MetadataFilter(key="event_name", value="Container Supply Chain Security Essentials", operator=FilterOperator.EQ)]),
                    similarity_top_k=10, llm=openai_llm)

qa_prompt = PromptTemplate(
    "Below given is all the required information about a course whose content i have loaded..\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "Given this content , I want you to create a detailed article on that course. You need to use all the given information to draft the article. "
    "You need to structure the article like heading, sub-heading, content, etc. It should be as elaborate as possible. I will provide you the name of the course as a query below."
    "he article should be written for a senior developer audience level. It should not be shown as a course. im just providing you the content , it should be a proper technical article. please dont include words like labs etc but all content must be used. "
    "Query: {query_str}\n"
    "Answer: "
)
class RAGStringQueryEngine(CustomQueryEngine):

    retriever: BaseRetriever
    response_synthesizer: BaseSynthesizer
    llm: LlamaOpenAI
    qa_prompt: PromptTemplate

    def custom_query(self, query_str: str):
        nodes = self.retriever.retrieve(query_str)

        context_str = "\n\n".join([n.node.get_content() for n in nodes])
        response = self.llm.complete(
            qa_prompt.format(context_str=context_str, query_str=query_str)
        )

        return str(response)
    
query_engine = RAGStringQueryEngine(
    retriever=retriever,
    response_synthesizer=synthesizer,
    llm=openai_llm,
    qa_prompt=qa_prompt,
)

response = query_engine.query("Please write an article for Container Supply Chain Security Essentials")

print(str(response))