"""
Chatbot with context and memory.
"""
import os

import openai
from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient
from azure.search.documents.models import Vector
from dotenv import load_dotenv

# Config for Azure Search.
AZURE_SEARCH_ENDPOINT = os.getenv("AZURE_SEARCH_ENDPOINT")
AZURE_SEARCH_KEY = os.getenv("AZURE_SEARCH_KEY")
AZURE_SEARCH_INDEX_NAME = "products-index-1"

# Config for Azure OpenAI.
AZURE_OPENAI_API_TYPE = "azure"
AZURE_OPENAI_API_BASE = os.getenv("AZURE_OPENAI_API_BASE")
AZURE_OPENAI_API_VERSION = "2023-03-15-preview"
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_CHATGPT_DEPLOYMENT = os.getenv("AZURE_OPENAI_CHATGPT_DEPLOYMENT")
AZURE_OPENAI_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")

# Chat roles
SYSTEM = "system"
USER = "user"
ASSISTANT = "assistant"


class Chatbot:
    """Chat with an LLM using RAG. Keeps chat history in memory."""

    chat_history = []

    def __init__(self):
        load_dotenv()
        openai.api_type = AZURE_OPENAI_API_TYPE
        openai.api_base = AZURE_OPENAI_API_BASE
        openai.api_version = AZURE_OPENAI_API_VERSION
        openai.api_key = AZURE_OPENAI_API_KEY

    def _summarize_user_intent(self, query: str) -> str:
        """
        Creates a user message containing the user intent, by summarizing the chat
        history and user query.
        """
        chat_history_str = ""
        for entry in self.chat_history:
            chat_history_str += f"{entry['role']}: {entry['content']}\n"
        messages = [
            {
                "role": SYSTEM,
                "content": (
                    "You're an AI assistant reading the transcript of a conversation "
                    "between a user and an assistant. Given the chat history and "
                    "user's query, infer user real intent."
                    f"Chat history: ```{chat_history_str}```\n"
                    f"User's query: ```{query}```\n"
                ),
            }
        ]
        chat_intent_completion = openai.ChatCompletion.create(
            deployment_id=AZURE_OPENAI_CHATGPT_DEPLOYMENT,
            messages=messages,
            temperature=0.7,
            max_tokens=1024,
            n=1,
        )
        user_intent = chat_intent_completion.choices[0].message.content

        return user_intent

    def _get_context(self, user_intent: str) -> list[str]:
        """
        Gets the relevant documents from Azure Cognitive Search.
        """
        query_vector = Vector(
            value=openai.Embedding.create(
                engine=AZURE_OPENAI_EMBEDDING_DEPLOYMENT, input=user_intent
            )["data"][0]["embedding"],
            fields="embedding",
        )

        search_client = SearchClient(
            endpoint=AZURE_SEARCH_ENDPOINT,
            index_name=AZURE_SEARCH_INDEX_NAME,
            credential=AzureKeyCredential(AZURE_SEARCH_KEY),
        )

        docs = search_client.search(search_text="", vectors=[query_vector], top=1)
        context_list = [doc["content"] for doc in docs]

        return context_list

    def _rag(self, context_list: list[str], query: str) -> str:
        """
        Asks the LLM to answer the user's query with the context provided.
        """
        user_message = {"role": USER, "content": query}
        self.chat_history.append(user_message)

        context = "\n\n".join(context_list)
        messages = [
            {
                "role": SYSTEM,
                "content": (
                    "You're a helpful assistant.\n"
                    "Please answer the user's question using only information you can "
                    "find in the context.\n"
                    "If the user's question is unrelated to the information in the "
                    "context, say you don't know.\n"
                    f"Context: ```{context}```\n"
                ),
            }
        ]
        messages = messages + self.chat_history

        chat_completion = openai.ChatCompletion.create(
            deployment_id=AZURE_OPENAI_CHATGPT_DEPLOYMENT,
            messages=messages,
            temperature=0.7,
            max_tokens=1024,
            n=1,
        )

        response = chat_completion.choices[0].message.content
        assistant_message = {"role": ASSISTANT, "content": response}
        self.chat_history.append(assistant_message)

        return response

    def ask(self, query: str) -> str:
        """
        Queries an LLM using RAG.
        """
        user_intent = self._summarize_user_intent(query)
        context_list = self._get_context(user_intent)
        response = self._rag(context_list, query)
        print(
            "*****\n"
            f"QUESTION:\n{query}\n"
            f"USER INTENT:\n{user_intent}\n"
            f"RESPONSE:\n{response}\n"
            "*****\n"
        )

        return response

src/1_openai/chatbot_1.py
bstollnitz/rag
"""
Initializes an Azure Cognitive Search index with our custom data, using vector search 
and semantic ranking.

To run this code, you must already have a "Cognitive Search" and an "OpenAI"
resource created in Azure.
"""
import os

import openai
from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    HnswParameters,
    HnswVectorSearchAlgorithmConfiguration,
    PrioritizedFields,
    SearchableField,
    SearchField,
    SearchFieldDataType,
    SearchIndex,
    SemanticConfiguration,
    SemanticField,
    SemanticSettings,
    SimpleField,
    VectorSearch,
)
from dotenv import load_dotenv
from langchain.document_loaders import DirectoryLoader, UnstructuredMarkdownLoader
from langchain.text_splitter import Language, RecursiveCharacterTextSplitter

# Config for Azure Search.
AZURE_SEARCH_ENDPOINT = os.getenv("AZURE_SEARCH_ENDPOINT")
AZURE_SEARCH_KEY = os.getenv("AZURE_SEARCH_KEY")
AZURE_SEARCH_INDEX_NAME = "products-index-1"

# Config for Azure OpenAI.
AZURE_OPENAI_API_TYPE = "azure"
AZURE_OPENAI_API_BASE = os.getenv("AZURE_OPENAI_API_BASE")
AZURE_OPENAI_API_VERSION = "2023-03-15-preview"
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")

DATA_DIR = "data/"


def load_and_split_documents() -> list[dict]:
    """
    Loads our documents from disc and split them into chunks.
    Returns a list of dictionaries.
    """
    # Load our data.
    loader = DirectoryLoader(
        DATA_DIR, loader_cls=UnstructuredMarkdownLoader, show_progress=True
    )
    docs = loader.load()

    # Split our documents.
    splitter = RecursiveCharacterTextSplitter.from_language(
        language=Language.MARKDOWN, chunk_size=6000, chunk_overlap=100
    )
    split_docs = splitter.split_documents(docs)

    # Convert our LangChain Documents to a list of dictionaries.
    final_docs = []
    for i, doc in enumerate(split_docs):
        doc_dict = {
            "id": str(i),
            "content": doc.page_content,
            "sourcefile": os.path.basename(doc.metadata["source"]),
        }
        final_docs.append(doc_dict)

    return final_docs


def get_index(name: str) -> SearchIndex:
    """
    Returns an Azure Cognitive Search index with the given name.
    """
    # The fields we want to index. The "embedding" field is a vector field that will
    # be used for vector search.
    fields = [
        SimpleField(name="id", type=SearchFieldDataType.String, key=True),
        SimpleField(name="sourcefile", type=SearchFieldDataType.String),
        SearchableField(name="content", type=SearchFieldDataType.String),
        SearchField(
            name="embedding",
            type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
            # Size of the vector created by the text-embedding-ada-002 model.
            vector_search_dimensions=1536,
            vector_search_configuration="default",
        ),
    ]

    # The "content" field should be prioritized for semantic ranking.
    semantic_settings = SemanticSettings(
        configurations=[
            SemanticConfiguration(
                name="default",
                prioritized_fields=PrioritizedFields(
                    title_field=None,
                    prioritized_content_fields=[SemanticField(field_name="content")],
                ),
            )
        ]
    )

    # For vector search, we want to use the HNSW (Hierarchical Navigable Small World)
    # algorithm (a type of approximate nearest neighbor search algorithm) with cosine
    # distance.
    vector_search = VectorSearch(
        algorithm_configurations=[
            HnswVectorSearchAlgorithmConfiguration(
                name="default",
                kind="hnsw",
                parameters=HnswParameters(metric="cosine"),
            )
        ]
    )

    # Create the search index.
    index = SearchIndex(
        name=name,
        fields=fields,
        semantic_settings=semantic_settings,
        vector_search=vector_search,
    )

    return index


def initialize(search_index_client: SearchIndexClient):
    """
    Initializes an Azure Cognitive Search index with our custom data, using vector
    search.
    """
    # Load our data.
    docs = load_and_split_documents()
    for doc in docs:
        doc["embedding"] = openai.Embedding.create(
            engine=AZURE_OPENAI_EMBEDDING_DEPLOYMENT, input=doc["content"]
        )["data"][0]["embedding"]

    # Create an Azure Cognitive Search index.
    index = get_index(AZURE_SEARCH_INDEX_NAME)
    search_index_client.create_or_update_index(index)

    # Upload our data to the index.
    search_client = SearchClient(
        endpoint=AZURE_SEARCH_ENDPOINT,
        index_name=AZURE_SEARCH_INDEX_NAME,
        credential=AzureKeyCredential(AZURE_SEARCH_KEY),
    )
    search_client.upload_documents(docs)


def delete(search_index_client: SearchIndexClient):
    """
    Deletes the Azure Cognitive Search index.
    """
    search_index_client.delete_index(AZURE_SEARCH_INDEX_NAME)


def main():
    load_dotenv()

    openai.api_type = AZURE_OPENAI_API_TYPE
    openai.api_base = AZURE_OPENAI_API_BASE
    openai.api_version = AZURE_OPENAI_API_VERSION
    openai.api_key = AZURE_OPENAI_API_KEY

    search_index_client = SearchIndexClient(
        AZURE_SEARCH_ENDPOINT, AzureKeyCredential(AZURE_SEARCH_KEY)
    )

    initialize(search_index_client)
    # delete(search_index_client)


if __name__ == "__main__":
    main()

src/1_openai/init_search_1.py
bstollnitz/rag
"""
Entry point for the chatbot.
"""
from chatbot_1 import Chatbot


def main():
    chatbot = Chatbot()
    chatbot.ask("I need a large backpack. Which one do you recommend?")
    chatbot.ask("How much does it cost?")
    chatbot.ask("And how much for a donut?")


if __name__ == "__main__":
    main()

src/1_openai/main_1.py
bstollnitz/rag
"""
Chatbot with context and memory, using LangChain.
"""
import os

from dotenv import load_dotenv
from langchain.chains import LLMChain
from langchain.chat_models import AzureChatOpenAI
from langchain.memory import ChatMessageHistory
from langchain.prompts.chat import ChatPromptTemplate, SystemMessagePromptTemplate
from langchain.retrievers.azure_cognitive_search import AzureCognitiveSearchRetriever

# Config for Azure OpenAI.
AZURE_OPENAI_API_TYPE = "azure"
AZURE_OPENAI_API_BASE = os.getenv("AZURE_OPENAI_API_BASE")
AZURE_OPENAI_API_VERSION = "2023-03-15-preview"
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_CHATGPT_DEPLOYMENT = os.getenv("AZURE_OPENAI_CHATGPT_DEPLOYMENT")

# Config for Azure Search.
AZURE_SEARCH_KEY = os.getenv("AZURE_SEARCH_KEY")
AZURE_SEARCH_SERVICE_NAME = os.getenv("AZURE_SEARCH_SERVICE_NAME")
AZURE_SEARCH_INDEX_NAME = "products-index-2"


class Chatbot:
    """Chat with an LLM using LangChain. Keeps chat history in memory."""

    history = ChatMessageHistory()

    def __init__(self):
        load_dotenv()

    def _summarize_user_intent(self, query: str) -> str:
        """
        Creates a user message containing the user intent, by summarizing the chat
        history and user query.
        """
        llm = AzureChatOpenAI(
            deployment_name=AZURE_OPENAI_CHATGPT_DEPLOYMENT,
            openai_api_type=AZURE_OPENAI_API_TYPE,
            openai_api_base=AZURE_OPENAI_API_BASE,
            openai_api_version=AZURE_OPENAI_API_VERSION,
            openai_api_key=AZURE_OPENAI_API_KEY,
            temperature=0.7,
        )

        chat_history_str = ""
        for entry in self.history.messages:
            chat_history_str += entry.type + ": " + entry.content + "\n "
        system_template = (
            "You're an AI assistant reading the transcript of a conversation "
            "between a user and an assistant. Given the chat history and "
            "user's query, infer user real intent."
            "Chat history: ```{chat_history_str}```\n"
            "User's query: ```{query}```\n"
        )
        chat_template = ChatPromptTemplate(
            messages=[SystemMessagePromptTemplate.from_template(system_template)]
        )

        llm_chain = LLMChain(llm=llm, prompt=chat_template)
        user_intent = llm_chain({"chat_history_str": chat_history_str, "query": query})[
            "text"
        ]

        return user_intent

    def _get_context(self, user_intent: str) -> list[str]:
        """
        Gets the relevant documents from Azure Cognitive Search.
        """
        retriever = AzureCognitiveSearchRetriever(
            api_key=AZURE_SEARCH_KEY,
            service_name=AZURE_SEARCH_SERVICE_NAME,
            index_name=AZURE_SEARCH_INDEX_NAME,
            top_k=3,
        )

        docs = retriever.get_relevant_documents(user_intent)
        context = [doc.page_content for doc in docs]

        return context

    def _rag(self, context_list: list[str], query: str) -> str:
        """
        Asks the LLM to answer the user's query with the context provided.
        """
        self.history.add_user_message(query)

        llm = AzureChatOpenAI(
            deployment_name=AZURE_OPENAI_CHATGPT_DEPLOYMENT,
            openai_api_type=AZURE_OPENAI_API_TYPE,
            openai_api_base=AZURE_OPENAI_API_BASE,
            openai_api_version=AZURE_OPENAI_API_VERSION,
            openai_api_key=AZURE_OPENAI_API_KEY,
            temperature=0.7,
        )

        context = "\n\n".join(context_list)
        system_template = (
            "You're a helpful assistant.\n"
            "Please answer the user's question using only information you can "
            "find in the context.\n"
            "If the user's question is unrelated to the information in the "
            "context, say you don't know.\n"
            "Context: ```{context}```\n"
        )
        chat_template = ChatPromptTemplate(
            messages=[SystemMessagePromptTemplate.from_template(system_template)]
            + self.history.messages
        )

        llm_chain = LLMChain(llm=llm, prompt=chat_template)
        response = llm_chain({"context": context})["text"]
        self.history.add_ai_message(response)

        return response

    def ask(self, query: str) -> str:
        """
        Queries an LLM using RAG.
        """
        user_intent = self._summarize_user_intent(query)
        context_list = self._get_context(user_intent)
        response = self._rag(context_list, query)
        print(
            "*****\n"
            f"QUESTION:\n{query}\n"
            f"USER INTENT:\n{user_intent}\n"
            f"RESPONSE:\n{response}\n"
            "*****\n"
        )

        return response

src/2_langchain/chatbot_2.py
bstollnitz/rag
"""
Initializes an Azure Cognitive Search index with our custom data, using vector search.
Uses LangChain.

To run this code, you must already have a "Cognitive Search" and an "OpenAI"
resource created in Azure.
"""
import os

from dotenv import load_dotenv
from langchain.document_loaders import DirectoryLoader, UnstructuredMarkdownLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
from langchain.vectorstores.azuresearch import AzureSearch
from langchain.vectorstores.utils import Document

# Config for Azure OpenAI.
AZURE_OPENAI_API_TYPE = "azure"
AZURE_OPENAI_API_BASE = os.getenv("AZURE_OPENAI_API_BASE")
AZURE_OPENAI_API_VERSION = "2023-03-15-preview"
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")

# Config for Azure Search.
AZURE_SEARCH_ENDPOINT = os.getenv("AZURE_SEARCH_ENDPOINT")
AZURE_SEARCH_KEY = os.getenv("AZURE_SEARCH_KEY")
AZURE_SEARCH_INDEX_NAME = "products-index-2"

DATA_DIR = "data/"


def load_and_split_documents() -> list[Document]:
    """
    Loads our documents from disc and split them into chunks.
    Returns a list of LancChain Documents.
    """
    # Load our data.
    loader = DirectoryLoader(
        DATA_DIR, loader_cls=UnstructuredMarkdownLoader, show_progress=True
    )
    docs = loader.load()

    # Split our documents.
    splitter = RecursiveCharacterTextSplitter.from_language(
        language=Language.MARKDOWN, chunk_size=6000, chunk_overlap=100
    )
    split_docs = splitter.split_documents(docs)

    return split_docs


def initialize():
    """
    Initializes an Azure Cognitive Search index with our custom data.
    """
    # Load our data.
    docs = load_and_split_documents()

    # Create an Azure Cognitive Search index.
    embeddings = OpenAIEmbeddings(
        deployment=AZURE_OPENAI_EMBEDDING_DEPLOYMENT,
        openai_api_base=AZURE_OPENAI_API_BASE,
        openai_api_key=AZURE_OPENAI_API_KEY,
        openai_api_version=AZURE_OPENAI_API_VERSION,
        openai_api_type=AZURE_OPENAI_API_TYPE,
    )
    vector_store = AzureSearch(
        azure_search_endpoint=AZURE_SEARCH_ENDPOINT,
        azure_search_key=AZURE_SEARCH_KEY,
        index_name=AZURE_SEARCH_INDEX_NAME,
        embedding_function=embeddings.embed_query,
    )

    # Upload our data to the index.
    vector_store.add_documents(documents=docs)


def main():
    load_dotenv()

    initialize()


if __name__ == "__main__":
    main()

src/2_langchain/init_search_2.py
bstollnitz/rag
"""
Entry point for the chatbot.
"""
from chatbot_2 import Chatbot


def main():
    chatbot = Chatbot()
    chatbot.ask("I need a large backpack. Which one do you recommend?")
    chatbot.ask("How much does it cost?")
    chatbot.ask("And how much for a donut?")


if __name__ == "__main__":
    main()

src/2_langchain/main_2.py
bstollnitz/rag
"""
Chatbot with context and memory, using Semantic Kernel.
"""
import os

import semantic_kernel as sk
from semantic_kernel.connectors.ai.open_ai import (
    AzureChatCompletion,
    OpenAITextEmbedding,
)
from semantic_kernel.connectors.memory.azure_cognitive_search import (
    AzureCognitiveSearchMemoryStore,
)
from semantic_kernel.semantic_functions.chat_prompt_template import ChatPromptTemplate
from semantic_kernel.semantic_functions.prompt_template_config import (
    PromptTemplateConfig,
)
from semantic_kernel.semantic_functions.semantic_function_config import (
    SemanticFunctionConfig,
)

# Config for Azure OpenAI.
AZURE_OPENAI_API_TYPE = "azure"
AZURE_OPENAI_API_BASE = os.getenv("AZURE_OPENAI_API_BASE")
AZURE_OPENAI_API_VERSION = "2023-03-15-preview"
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_CHATGPT_DEPLOYMENT = os.getenv("AZURE_OPENAI_CHATGPT_DEPLOYMENT")
AZURE_OPENAI_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")

# Config for Azure Search.
AZURE_SEARCH_ENDPOINT = os.getenv("AZURE_SEARCH_ENDPOINT")
AZURE_SEARCH_KEY = os.getenv("AZURE_SEARCH_KEY")
AZURE_SEARCH_INDEX_NAME = "products-index-3"

# Chat roles
SYSTEM = "system"
USER = "user"
ASSISTANT = "assistant"

PLUGIN_NAME = "rag_plugin"


class Chatbot:
    """Chat with an LLM. Keeps chat history in memory."""

    kernel = None
    variables = None

    def __init__(self):
        # Create a kernel and adds the Azure OpenAI connector to it.
        self.kernel = sk.Kernel()
        self.kernel.add_chat_service(
            "azureopenai",
            AzureChatCompletion(
                deployment_name=AZURE_OPENAI_CHATGPT_DEPLOYMENT,
                endpoint=AZURE_OPENAI_API_BASE,
                api_key=AZURE_OPENAI_API_KEY,
            ),
        )

        # Create the variables to pass to the chat functions.
        self.variables = sk.ContextVariables()
        self.variables["chat_history"] = ""

    async def _summarize_user_intent(self, query: str) -> str:
        """
        Creates a user message containing the user intent, by summarizing the chat
        history and user query.
        """
        # Define the user template.
        self.variables["query"] = query
        user_template = (
            "You're an AI assistant reading the transcript of a conversation "
            "between a user and an assistant. Given the chat history and "
            "user's query, infer user real intent."
            "Chat history: ```{{$chat_history}}```\n"
            "User's query: ```{{$query}}```\n"
        )

        # Create a semantic function.
        prompt_config_dict = {
            "type": "completion",
            "description": "An AI assistant that infers user intent.",
            "completion": {
                "temperature": 0.7,
                "top_p": 0.5,
                "max_tokens": 200,
                "number_of_responses": 1,
                "presence_penalty": 0,
                "frequency_penalty": 0,
            },
            "input": {
                "parameters": [
                    {
                        "name": "query",
                        "description": "The question asked by the user.",
                        "defaultValue": "",
                    },
                    {
                        "name": "chat_history",
                        "description": "All the user and assistant messages so far .",
                        "defaultValue": "",
                    },
                ]
            },
        }
        prompt_config = PromptTemplateConfig.from_dict(prompt_config_dict)
        prompt_template = ChatPromptTemplate(
            template=user_template,
            prompt_config=prompt_config,
            template_engine=self.kernel.prompt_template_engine,
        )
        user_intent_function_config = SemanticFunctionConfig(
            prompt_config, prompt_template
        )
        user_intent_function = self.kernel.register_semantic_function(
            skill_name=PLUGIN_NAME,
            function_name="user_intent_function",
            function_config=user_intent_function_config,
        )

        # Run the semantic function.
        response = await self.kernel.run_async(
            user_intent_function, input_vars=self.variables
        )

        return str(response)

    async def _get_context(self, query: str) -> list[str]:
        """
        Gets the relevant documents from Azure Cognitive Search.
        """
        kernel = sk.Kernel()
        kernel.add_text_embedding_generation_service(
            "openai-embedding",
            OpenAITextEmbedding(
                model_id=AZURE_OPENAI_EMBEDDING_DEPLOYMENT,
                api_key=AZURE_OPENAI_API_KEY,
                endpoint=AZURE_OPENAI_API_BASE,
                api_type=AZURE_OPENAI_API_TYPE,
                api_version=AZURE_OPENAI_API_VERSION,
            ),
        )
        kernel.register_memory_store(
            memory_store=AzureCognitiveSearchMemoryStore(
                vector_size=1536,
                search_endpoint=AZURE_SEARCH_ENDPOINT,
                admin_key=AZURE_SEARCH_KEY,
            )
        )

        docs = await kernel.memory.search_async(AZURE_SEARCH_INDEX_NAME, query, limit=1)
        context = [doc.text for doc in docs]

        return context

    async def _rag(self, context_list: list[str], query: str) -> str:
        """
        Asks the LLM to answer the user's query with the context provided.
        """
        # Define the system template.
        context = "\n\n".join(context_list)
        self.variables["context"] = context
        system_template = (
            "You're a helpful assistant.\n"
            "Please answer the user's question using only information you can "
            "find in the context.\n"
            "If the user's question is unrelated to the information in the "
            "context, say you don't know.\n"
            "Context: ```{{$context}}```\n"
        )

        # Define the user template.
        self.variables["query"] = query
        user_template = "{{$chat_history}}" + f"{USER}: " + "{{$query}}\n"

        # Create a semantic function.
        prompt_config_dict = {
            "type": "completion",
            "description": "A chatbot that's a helpful assistant.",
            "completion": {
                "temperature": 0.7,
                "top_p": 0.5,
                "max_tokens": 200,
                "number_of_responses": 1,
                "chat_system_prompt": system_template,
                "presence_penalty": 0,
                "frequency_penalty": 0,
            },
            "input": {
                "parameters": [
                    {
                        "name": "query",
                        "description": "The question asked by the user.",
                        "defaultValue": "",
                    },
                    {
                        "name": "context",
                        "description": "The context for the assistant's answer.",
                        "defaultValue": "",
                    },
                    {
                        "name": "chat_history",
                        "description": "All the user and assistant messages so far.",
                        "defaultValue": "",
                    },
                ]
            },
        }
        prompt_config = PromptTemplateConfig.from_dict(prompt_config_dict)
        prompt_template = ChatPromptTemplate(
            template=user_template,
            prompt_config=prompt_config,
            template_engine=self.kernel.prompt_template_engine,
        )
        rag_function_config = SemanticFunctionConfig(prompt_config, prompt_template)
        rag_function = self.kernel.register_semantic_function(
            skill_name=PLUGIN_NAME,
            function_name="rag_function",
            function_config=rag_function_config,
        )

        # Run the semantic function.
        response = await self.kernel.run_async(rag_function, input_vars=self.variables)
        self.variables["chat_history"] += f"{USER}: {query}\n{ASSISTANT}: {response}\n"

        return str(response)

    async def ask(self, query: str) -> str:
        """
        Queries an LLM using RAG.
        """
        user_intent = await self._summarize_user_intent(query)
        context_list = await self._get_context(user_intent)
        response = await self._rag(context_list, query)
        print(
            "*****\n"
            f"QUESTION:\n{query}\n"
            f"USER INTENT:\n{user_intent}\n"
            f"RESPONSE:\n{response}\n"
            "*****\n"
        )

        return response

src/3_semantic_kernel/chatbot_3.py
bstollnitz/rag
"""
Initializes an Azure Cognitive Search index with our custom data, using vector search.
Uses Semantic Kernel.

To run this code, you must already have a "Cognitive Search" and an "OpenAI"
resource created in Azure.
"""
import asyncio
import ntpath
import os

import semantic_kernel as sk
from dotenv import load_dotenv
from langchain.document_loaders import DirectoryLoader, UnstructuredMarkdownLoader
from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
from semantic_kernel.connectors.ai.open_ai import OpenAITextEmbedding
from semantic_kernel.connectors.memory.azure_cognitive_search import (
    AzureCognitiveSearchMemoryStore,
)

# Config for Azure OpenAI.
AZURE_OPENAI_API_TYPE = "azure"
AZURE_OPENAI_API_BASE = os.getenv("AZURE_OPENAI_API_BASE")
AZURE_OPENAI_API_VERSION = "2023-03-15-preview"
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")

# Config for Azure Search.
AZURE_SEARCH_ENDPOINT = os.getenv("AZURE_SEARCH_ENDPOINT")
AZURE_SEARCH_KEY = os.getenv("AZURE_SEARCH_KEY")
AZURE_SEARCH_INDEX_NAME = "products-index-3"


DATA_DIR = "data/"


def load_and_split_documents() -> list[dict]:
    """
    Loads our documents from disc and split them into chunks.
    Returns a list of dictionaries.
    """
    # Load our data.
    loader = DirectoryLoader(
        DATA_DIR, loader_cls=UnstructuredMarkdownLoader, show_progress=True
    )
    docs = loader.load()

    # Split our documents.
    splitter = RecursiveCharacterTextSplitter.from_language(
        language=Language.MARKDOWN, chunk_size=6000, chunk_overlap=100
    )
    split_docs = splitter.split_documents(docs)

    # Convert our LangChain Documents to a list of dictionaries.
    final_docs = []
    for i, doc in enumerate(split_docs):
        doc_dict = {
            "id": str(i),
            "content": doc.page_content,
            "sourcefile": ntpath.basename(doc.metadata["source"]),
        }
        final_docs.append(doc_dict)

    return final_docs


async def initialize(memory_store: AzureCognitiveSearchMemoryStore):
    """
    Initializes an Azure Cognitive Search index with our custom data.
    """
    # Load our data.
    docs = load_and_split_documents()

    # Create an Azure Cognitive Search index.
    kernel = sk.Kernel()
    kernel.add_text_embedding_generation_service(
        "openai-embedding",
        OpenAITextEmbedding(
            model_id=AZURE_OPENAI_EMBEDDING_DEPLOYMENT,
            api_key=AZURE_OPENAI_API_KEY,
            endpoint=AZURE_OPENAI_API_BASE,
            api_type=AZURE_OPENAI_API_TYPE,
            api_version=AZURE_OPENAI_API_VERSION,
        ),
    )
    kernel.register_memory_store(memory_store)

    # Upload our data to the index.
    for doc in docs:
        await kernel.memory.save_information_async(
            AZURE_SEARCH_INDEX_NAME, id=doc["id"], text=doc["content"]
        )


async def delete(memory_store: AzureCognitiveSearchMemoryStore):
    """
    Deletes the Azure Cognitive Search index.
    """
    await memory_store.delete_collection_async(AZURE_SEARCH_INDEX_NAME)


async def main():
    load_dotenv()

    memory_store = AzureCognitiveSearchMemoryStore(
        vector_size=1536,
        search_endpoint=AZURE_SEARCH_ENDPOINT,
        admin_key=AZURE_SEARCH_KEY,
    )

    await initialize(memory_store)
    # await delete(memory_store)


if __name__ == "__main__":
    asyncio.run(main())

src/3_semantic_kernel/init_search_3.py
bstollnitz/rag
"""
Entry point for the chatbot.
"""
import asyncio

from chatbot_3 import Chatbot


async def main():
    chatbot = Chatbot()
    await chatbot.ask("I need a large backpack. Which one do you recommend?")
    await chatbot.ask("How much does it cost?")
    await chatbot.ask("And how much for a donut?")


if __name__ == "__main__":
    asyncio.run(main())

src/3_semantic_kernel/main_3.py
bstollnitz/rag