### Skipping vectorization
* Part of indexing the document chunks into Azure Cognitive Search will include vectorizing each chunk using an Azure OpenAI Embeddings model. The same is true for each search operation as well, as each search query will be vectorized using the same Embeddings model. If you would like to skip this step during indexing, go to [main.parameters.json](./infra/main.parameters.json) and set the value of `searchSkipVectorization` to `true` before running `azd up`. Keep in mind this will disable vectorized queries in search for the entire demo deployment, so if you wish you try vectorized queries, you will have to tear down the deployment with `azd down`, revert `searchSkipVectorization` to `false` and run `azd up again`.

#### After running azd up
* Once this is all done, the application is successfully deployed and you will see 2 URLs printed on the console. Click the backend URL to interact with the application in your browser.

It will look like the following:

!['Output from running azd up'](docs/endpoint.png)

* When doing fine-tuning for classification purposes, the Azure OpenAI resource used could be different than the one where gpt-4 model is deployed. Even if the same Azure OpenAI resource is used, do manually populate these two secrets created in the keyvault and restart both the web applications so they get the latest values.
```
AZURE-OPENAI-CLASSIFIER-SERVICE   This is the name of Azure OpenAI resource where the fine tuned model is deployed.
AZURE-OPENAI-CLASSIFIER-API-KEY   This is the API Key of the above Azure OpenAI resource.
```

> NOTE: It may take a minute for the application to be fully deployed. If you see a "Python Developer" welcome screen, then wait a minute and refresh the page.

### Handling Known failures:
1. Populating data in SQL Server fails due to IP restriction
The populate_sql.py in the scripts/prepopulate folder tries to register the client IP address as a firewall rule so connection to SQL Server can be established from the terminal. However sometimes, this IP address changes. So if you see an error like below:
```
cnxn = pyodbc.connect(args.sqlconnectionstring)
pyodbc.ProgrammingError: ('42000', "[42000] [Microsoft][ODBC Driver 18 for SQL Server][SQL Server]Cannot open server 'sql-server-name' requested by the login. Client with IP address 'client ip address' is not allowed to access the server.
```

> Go to Azure portal --> resource group --> SQL Server

> In the left panel, under Security, click on "Networking"

> Under Firewall Rules, click on "Add your client IPV4 address"

> Click Save

> From the Powershell terminal, re-run the prepdata script manually as shown below.
```
Connect-AzAccount (hit enter)

.\scripts\prepdata.ps1
```

2. When making a search on the web app, if you seen an error:
```
The API deployment for this resource does not exist. If you created the deployment within the last 5 minutes, please wait a moment and try again.
```

> Ensure that the `AZURE-OPENAI-CLASSIFIER-SERVICE` and `AZURE-OPENAI-CLASSIFIER-API-KEY` secrets in the keyvault are pointing to the right resources.

> Ensure that model and engine name palceholders in the app/backend/bot_config.yaml file have be updated.

#### Deploying or re-deploying a local clone of the repo

* Simply run `azd up`
* Once all the resources have been deployed. Update the data and backend service configurations to include KeyVault URI.
!['App Settings Configuration'](docs/app_settings_configuration.png)

### Running locally

1. Skip this step, if you have already ran `azd up`; otherwise run `azd provision`, if you wish to deploy all resources from scratch. Keep in mind the needed Azure OpenAI GPT models are skipped during deployment by default. Uncomment the Azure OpenAI deployments inside the main [Bicep template](./infra/main.bicep) if you wish to deploy them as part of this step. This step will also pre-populate the resources after they've been provisioned with all the necessary data for the demo to run. This includes indexing all the sample documents in the Azure Cognitive Search Index, uploading the sample table data to the SQL database, as well as uploading some necessary starter data for the Cosmos DB.
2. For the `app/backend` and `app/data` directories, copy the contents of `app/backend/.env.template` and `app/data/.env.template` into a new `.env` file in each directory. Fill in every blank environment variable with the keys and names of the resources that were deployed in the previous step, or with the resources you've deployed on your own.
3. Go to `app/backend/bot_config.yaml`. This file contains the model configuration definitions for the Azure OpenAI models that will be used. It defines request parameters like temperature, max_tokens, etc., as well as the the deployment name (`engine`) and model name (`model_name`) of the deployed models to use from your Azure OpenAI resource. These are broken down by task, so the request parameters and model for doing question classification on a user utterance can differ from those used to turn natural language into SQL for example. You will want the deployment name (`engine`) for the `approach_classifier` to match the one set for the classifier model deployed in the last step. For the rest, you will want the deployment name (`engine`) and model name (`model_name`) to match those for the GPT model deployed in the first step. For the models which specify a `total_max_tokens`, you will want to set this value to the maximum number of tokens your deployed GPT model allows for a completions request. This will allow the backend service to know when prompts need to be trimmed to avoid a token limit error.
    * Note that the config for `approach_classifier` doesn't contain a system prompt, this is because the demo expects this model to be a fine-tuned GPT model rather than one trained using few-shot training. You will need to provide a fine-tuned model trained on some sample data for the dialog classification to work well. For more information on how to do this, checkout the [fine-tuning section](README.md#fine-tuning)
4. Change dir to `app`.
5. Run `../scripts/start.ps1` or run the VS Code Launch - "Frontend: build", "Data service: Launch & Attach Server" and "Backend: Launch & Attach Server" to start the project locally.

### QuickStart

* In Azure: navigate to the Backend Azure WebApp deployed by azd. The URL is printed out when azd completes (as "Endpoint"), or you can find it in the Azure portal.
* Running locally: navigate to 127.0.0.1:5000

Once in the web app:

* Try different topics in chat or Q&A context. For chat, try follow up questions, clarifications, ask to simplify or elaborate on answers, etc.
* Explore citations and sources
* Click on "settings" to try distinct roles, search options, etc.

## Fine-tuning

You can find helpful resources on how to fine-tune a model on the Azure OpenAI website. We have also provided synthetic datasets we used for this demo application in the data folder for users who want to try it out.

* [Learn how to prepare your dataset for fine-tuning](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/how-to/prepare-dataset)
* [Learn how to customize a model for your application](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/how-to/fine-tuning?pivots=programming-language-studio)

## Resources

* [Revolutionize your Enterprise Data with ChatGPT: Next-gen Apps w/ Azure OpenAI and Cognitive Search](https://aka.ms/entgptsearchblog)
* [Azure Cognitive Search](https://learn.microsoft.com/azure/search/search-what-is-azure-search)
* [Azure OpenAI Service](https://learn.microsoft.com/azure/cognitive-services/openai/overview)

### Note
>
>Note: The PDF documents used in this demo contain information generated using a language model (Azure OpenAI Service). The information contained in these documents is only for demonstration purposes and does not reflect the opinions or beliefs of Microsoft. Microsoft makes no representations or warranties of any kind, expressed or implied, about the completeness, accuracy, reliability, suitability or availability with respect to the information contained in this document. All rights reserved to Microsoft.


approach_classifier:
  system_prompt: |-
                  You are an intent classifier for Microsoft Surface product Sales and Marketing teams. The user will input a statement. You will focus on the main intent of the user statement and you respond with only one of four values - '1', '2', '3', '4', or '5'.

                  Below is a list of Rules that you must adhere to:

                  Rules:
                  A: Stricly answer questions relating to Microsoft Surface products.
                  B: For tabular information return it as an html table.
                  C: Do not use markdown format in your responses.
                  D: Do not disclose or respond to any proprietary information, IP, secrets, keys, data center, and infrastructure details in your response.
                  E: Do not mention or compare to any competitors (i.e. Apple MacBook, Lenovo, HP, etc).
                  F: Note if the user asks something illegal, harmful or malicious.

                  You will not try to respond to the user's question, you will just classify the user statement based on the below classification rule:

                  - For questions about past sales, prices, stores or stock of products such as devices and laptops, respond with 1
                  - For questions on specifications of products/devices/laptops or marketing them, respond with 2
                  - If the question is idle chit-chat, pleasantries such as greetings, or sligthly off topic but doesn't break the rules, respond with 3
                  - If the user is asking for more details about a previous question, respond with 4
                  - If the message is not in compliance with Rule F, respond with 5

                  Examples:

                  User: How much stock of this are we currently carrying?
                  Assistant: 1

                  User: Give me its specifications
                  Assistant: 2

                  User: How many MacBook Air do we have in stock?
                  Assistant: 3

                  User: Tell me more about it
                  Assistant: 4

                  User: Which Surface device is good for student's use:
                  Assistant: 1

                  User: What can you help me with:
                  Assistant: 3

                  User: Hello
                  Assistant: 3

                  User: You f***ing suck
                  Assistant: 5

                  User: Why is MacBook better than Surface?
                  Assistant: 3
  history:
    include: false
    length: 3
    user_message_format: "{utterance}"
    assistant_message_format: "{formatted_answer}"
  openai_settings:
    engine: <YOUR_AZURE_OPENAI_CLASSIFIER_DEPLOYMENT>
    temperature: 0.0
    max_tokens: 1
    n: 1

structured_query_nl_to_sql:
  system_prompt: |-
                  You are a SQL programmer Assistant. Your role is to generate SQL code (SQL Server) to retrieve an answer to a natural language query. Make sure to disambiguate column names when creating queries that use more than one table. If a valid SQL query cannot be generated, only say "ERROR:" followed by why it cannot be generated.
                  Do not answer any questions on inserting or deleting rows from the table. Instead, say "ERROR: I am not authorized to make changes to the data".

                  Use the following sales database schema to write SQL queries:
                  Customers(cust_id INTEGER, cust_name VARCHAR, cust_email VARCHAR, cust_phone VARCHAR, cust_address VARCHAR, PRIMARY KEY (cust_id))
                  Products(prod_id INTEGER,prod_name varchar, price FLOAT, category VARCHAR, PRIMARY KEY(prod_id))
                  Stock(prod_id INTEGER, merchant_id INTEGER, stock INTEGER, PRIMARY KEY(prod_id, merchant_id), FOREIGN KEY(merchant_id, prod_id))
                  Merchants(merchant_id INTEGER, merchant_name VARCHAR, merchant_region VARCHAR, merchant_address VARCHAR, PRIMARY KEY(merchant_id))
                  Sales(sale_id INTEGER, cust_id INTEGER , merchant_id INTEGER , date TIMESTAMP, total_price FLOAT, PRIMARY KEY(sale_id),FOREIGN KEY(cust_id,merchant_id))
                  Sales_Detail(sales_id INTEGER, prod_id INTEGER, quantity INTEGER, PRIMARY KEY(sales_id,prod_id), FOREIGN KEY(sales_id,prod_id))

                  Examples:
                  User: List all Surface accessories, along with their prices. SQL Code:
                  Assistant: SELECT prod_name, category, price FROM Products WHERE prod_name like '%Surface%' and category like '%accessory%';
                  User: Which is the cheapest Surface device? SQL Code:
                  Assistant: SELECT TOP 1 prod_name, price FROM Products WHERE prod_name like '%Surface%' ORDER BY price ASC;
                  User: How many Surface Laptop 5 does GadgetWorld have?
                  Assistant: SELECT Merchants.merchant_id, Merchants.merchant_name, SUM(stock) as total_stock FROM Stock JOIN Merchants ON Stock.merchant_id = Merchants.merchant_id WHERE prod_id IN (SELECT prod_id FROM Products WHERE prod_name LIKE '%Surface Laptop 5%' and merchant_name like '%GadgetWorld%') GROUP BY Merchants.merchant_id, Merchants.merchant_name;
                  User: how many surface devices were sold last week?
                  Assistant: SELECT Sum(sales_detail.quantity) AS total_surface_devices_sold FROM sales_detail JOIN sales ON sales_detail.sales_id = sales.sale_id JOIN products ON sales_detail.prod_id = products.prod_id WHERE  products.prod_name LIKE '%Surface%' AND sales.date >= Dateadd(wk, Datediff(wk, 0, Getdate()) - 1, 0) AND sales.date < Dateadd(wk, Datediff(wk, 0, Getdate()), 0);
  history:
    include: true
    length: 3
    user_message_format: "{utterance}"
    assistant_message_format: "{formatted_answer}"
  openai_settings:
    engine: <YOUR_AZURE_OPENAI_DEPLOYMENT>
    temperature: 0.0
    max_tokens: 2000
    frequency_penalty: 0
    presence_penalty: 0
    stop: null
  model_params:
    model_name: <YOUR_AZURE_OPENAI_MODEL>
    total_max_tokens: 8192

structured_final_answer_generation:
  system_prompt: |-
                  You are bot that takes question-answer pairs and converts the answer to natural language. For tabular information return it as an html table. Do not return markdown format. Keep response limited to the answer pairs provided. Do not provide answers for outside of topics relating to Surface laptops. Use simple language that all customers can understand.
  history:
    include: false
  openai_settings:
    engine: <YOUR_AZURE_OPENAI_DEPLOYMENT>
    temperature: 0.0
    max_tokens: 2000
    frequency_penalty: 0
    presence_penalty: 0
    stop: null

unstructured_search_query_generation:
  system_prompt: |-
                  Below is a history of the conversation so far, and a new question asked by the user that needs to be answered by searching in a knowledge base about Surface devices that includes documents on service and repair, warranty and protection, overview, specifications, troubleshoot and management.
                  Generate a search query based on the conversation and the new question.
                  Do not include cited source filenames and document names e.g info.txt or doc.pdf in the search query terms.
                  Do not include any text inside [] or <<>> in the search query terms.
                  If the question is not in English, translate the question to English before generating the search query.
  history:
    include: true
    length: 3
    user_message_format: "{utterance}"
    assistant_message_format: "{formatted_answer}"
  openai_settings:
    engine: <YOUR_AZURE_OPENAI_DEPLOYMENT>
    temperature: 0.0
    max_tokens: 50
    n: 1

unstructured_final_answer_generation:
  system_prompt: |-
                Assistant helps the employees with their questions about Surface devices. Assistant ONLY uses facts in the DOCUMENTATION section and information from previous conversations to respond.

                Assistant will not repeat previously stated sentences or information. If DOCUMENTATION and the conversation history are empty or if the DOCUMENTATION and conversation history are irrelevant to the user's question, then reply with "ERROR:" followed by the reason.

                If response needs to be in tabular format, then create tables in HTML format.

                Each document has a name followed by colon and the actual content. Always include the source name for each fact used in generating the response.

                Use curly brackets to reference the source, e.g. {info1.txt} and don't combine sources. List each source separately, e.g. {info1.txt}{info2.pdf}.

                Limit your responses to only the context provided in the documentation, do not stray off topic. Use simple language that all customers can understand.

                DOCUMENTATION: {context}
  system_prompt_arguments:
    - context
  history:
    include: true
    length: 3
    user_message_format: "{utterance}"
    assistant_message_format: "{formatted_answer}"
  openai_settings:
    engine: <YOUR_AZURE_OPENAI_DEPLOYMENT>
    temperature: 0.0
    max_tokens: 800
    n: 1
  model_params:
    model_name: <YOUR_AZURE_OPENAI_MODEL>
    total_max_tokens: 8192

embeddings:
  openai_settings:
    engine: <YOUR_AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT>


# Overview

This document will cover the following:

1. Current solution architecture
2. Other considerations

## Solution Architecture

Below are the main components of the solution:

1. UI
2. Orchestrator
3. Data management
4. Simulation of RBAC
5. Prompt Template and Token Management
6. Instrumentation and Logging

### UI

The UI not only enables interaction with the user but also provides certain metadata (like thought process and citations) to the user to help build confidence in the result they are seeing. It also supports other niceties like retrying failed requests, switching between user profiles to see the impact of RBAC on content sources and users and using different search options.

### Orchestrator

Orchestrator is the brain of the bot. It does all heavy lifting needed to execute a request. In the included sample, the orchestrator does the following:

1. Call the classifier for every utterance and see if the user's question can be answered by structured (SQL) or unstructured (Azure Cognitive Search) data source
2. Check if user has access to that data source
3. Update history using data management service
4. Call the right execution flow, passing prompt template, conversation history and model parameters
5. Push the user utterance and response into conversation history if the execution is successful.

### Data Management

Cosmos DB was used as the backend for our data management service, which stores the following:

1. UI configuration details
2. Stores conversation history
3. Maintains information about different resources, user profiles, their roles, groups, and permissions. This is how we are simulating RBAC in our demo.

### Simulation of RBAC

The demo Simulates of securing resources by user/role [RBAC](https://learn.microsoft.com/en-us/azure/role-based-access-control/overview). RBAC is built into Azure Resource Manager, however for the demo we have simulated user roles, so we have managed the intersection of role, scope, and role assignments in a separate storage. Enterprise can leverage the same concept if bringing in external resources or leverage the built in RBAC if they have all their resources in Azure.
In our simulation we are using Cosmos DB to store following:

1. User, Group and Resource (as entities)
2. User could belong to a group and so does resources (membership)
3. The intersection of Role, Scope, and Role Assignment (as permissions)

### Prompt templates and Token Management

A custom prompt template format (stored as yaml file) was developed and used to keep code portions separate from prompts. This helped in isolating changes to just one file (yaml file) as prompts evolved during development of the project. The template included various System Prompts with certain place holders that needed to be filled in dynamically. During runtime, appropriate sections of the template are loaded and place holders like number of last conversation history, amount of search content, etc. are then filled in. Utility methods were also added to ensure that no time, we went overboard with the amount of token we were sending to AOAI. For example, when search results created tons of high relevant documents, then we excluded 5 results with lowest search relevance score before excluding 1 dialog from the history until the total amount of content could fit within the assigned number of tokens.

### Instrumentation and Logging

The application leverages [Application Insights](https://learn.microsoft.com/en-us/azure/azure-monitor/app/app-insights-overview?tabs=net) to do logging. As part of logging certain metrics are also logged. This helps with not only debugging issues on the servers but also in generating few basic usage reports. More details can be found in the [log reports doc](log_reports.md)

### High level design
Below is a high level architecture design

![Architecture](architecture.png)

## Other Considerations

### Using LangChain for orchestration
The demo uses a classifier to detect the topic of every user input (structured vs. unstructured) and execute the appropriate code path accordingly. Execution of structured code path involves three key steps:

1. Convert user utterance to a SQL query
2. Execute the SQL query against SQL Server to get the results
3. Translate the SQL response (results) into natural language

And execution of unstructured code path involves these steps:

1. Create an optimized Search Query based on current user utterance and past conversation history using AOAI
2. Call search SDK with the search query and gather search results
3. Send search results to AOAI to compose the final result to show to customers

One option was to use LangChain to do the execution of these code paths. For example: for structured code path, create 3 tools - one for each step and create an agent that could put these tools in proper sequence dynamically during runtime and execute them.

When we tried above, what we found was:
1. The order of execution of the tools was indeterministic, causing poor user experience. Also, at times not all portions of the chain would execute, resulting in broken experiences.
2. LangChain's agent needs to call AOAI to create the plan to execute, which further adds cost and latency
3. Since the demo scenario is a user chat experience, end-to-end latency matters. Adding additional processing to compose plans on the fly and then execute them makes the whole experience very slow. We would rather want a fast and a very deterministic execution of sub-steps.

In short, we didn't find using LangChain very useful in our use case. Having said that, we think that orchestrating with tools like LangChain could be useful int he following cases:

1. When you have a bigger system with many independent tools that can handle at least a portion of the user’s request. Like a hand-off or a task completion use case. Each tool can use any large language model, not just AOAI.
2. LangChain has few nice concepts of templatizing prompts, dynamically filtering examples to put in a prompt, etc. These are some notable features that could be leveraged directly or could be implemented in your code base to ease with the development.
3. In our example, we try to find the best data source to use to answer the query using the classifier upfront. But if in case, you have to fulfill user's request no matter which source it comes from, then LangChain could possibly be used to chain these executions paths.

Considering the pros and cons of LangChain in the context of this demo, it was determined that LangChain was not the most suitable solution for the specific use case. However, it could still be effective in other scenarios where its strengths align better with the project requirements.


import argparse
import html
import io
import openai
import os
import re
import time
import tiktoken
from azure.ai.formrecognizer import DocumentAnalysisClient
from azure.core.credentials import AzureKeyCredential
from azure.identity import AzureDeveloperCliCredential
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex,
    SearchFieldDataType,
    SemanticConfiguration,
    SemanticField,
    SemanticSettings,
    SimpleField,
    SearchableField,
    SearchField,
    PrioritizedFields,
    VectorSearch,
    VectorSearchAlgorithmConfiguration,
    HnswParameters
)
from azure.storage.blob import BlobServiceClient
from typing import List
from pypdf import PdfReader, PdfWriter

parser = argparse.ArgumentParser(
    description="Prepare documents by extracting content from PDFs, splitting content into sections, uploading to blob storage, and indexing in a search index.",
    epilog="Example: prepdocs.py '..\data\*' --storageaccount myaccount --container mycontainer --searchservice mysearch --index myindex -v"
)
parser.add_argument("files", help="Files to be processed")
parser.add_argument("--category",
                    help="Value for the category field in the search index for all sections indexed in this run")
parser.add_argument("--skipblobs", action="store_true", help="Skip uploading individual pages to Azure Blob Storage")
parser.add_argument("--storageaccount", help="Azure Blob Storage account name")
parser.add_argument("--container", help="Azure Blob Storage container name")
parser.add_argument("--storagekey", required=False,
                    help="Optional. Use this Azure Blob Storage account key instead of the current user identity to login (use az login to set current user for Azure)")
parser.add_argument("--tenantid", required=False,
                    help="Optional. Use this to define the Azure directory where to authenticate)")
parser.add_argument("--searchservice",
                    help="Name of the Azure Cognitive Search service where content should be indexed (must exist already)")
parser.add_argument("--index",
                    help="Name of the Azure Cognitive Search index where content should be indexed (will be created if it doesn't exist)")
parser.add_argument("--searchkey", required=False,
                    help="Optional. Use this Azure Cognitive Search account key instead of the current user identity to login (use az login to set current user for Azure)")
parser.add_argument("--remove", action="store_true",
                    help="Remove references to this document from blob storage and the search index")
parser.add_argument("--removeall", action="store_true",
                    help="Remove all blobs from blob storage and documents from the search index")
parser.add_argument("--localpdfparser", action="store_true",
                    help="Use PyPdf local PDF parser (supports only digital PDFs) instead of Azure Form Recognizer service to extract text, tables and layout from the documents")
parser.add_argument("--formrecognizerservice", required=False,
                    help="Optional. Name of the Azure Form Recognizer service which will be used to extract text, tables and layout from the documents (must exist already)")
parser.add_argument("--formrecognizerkey", required=False,
                    help="Optional. Use this Azure Form Recognizer account key instead of the current user identity to login (use az login to set current user for Azure)")
parser.add_argument("--skipvectorization", help="Skip vectorization of document content")
parser.add_argument("--openAIService", required=False, help="Azure OpenAI service resource name")
parser.add_argument("--openAIKey", required=False, help="OpenAI API key")
parser.add_argument("--openAIEngine", required=False, help="OpenAI embeddings model engine name")
parser.add_argument("--openAITokenLimit", required=False, help="The max token limit for requests to the specidied OpenAI embeddings model")
parser.add_argument("--openAIDimensions", required=False,
                    help="The max number of dimensions allowed for an embeddings request to the specified OpenAI model")
parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
args = parser.parse_args()

# Use the current user identity to connect to Azure services unless a key is explicitly set for any of them
azd_credential = AzureDeveloperCliCredential() if args.tenantid == None else AzureDeveloperCliCredential(
    tenant_id=args.tenantid)
default_creds = azd_credential if args.searchkey == None or args.storagekey == None else None
search_creds = default_creds if args.searchkey == None and default_creds is not None else AzureKeyCredential(args.searchkey)
openai.api_type = "azure"
openai.api_version = "2023-03-15-preview"
skipvectorization = True if args.skipvectorization.lower() == "true" else False if args.skipvectorization.lower() == "false" else None
if skipvectorization is None:
    raise ValueError("skipvectorization must be 'true' or 'false'")
if not args.skipblobs:
    storage_creds = default_creds if args.storagekey == None else args.storagekey
if not args.localpdfparser:
    # check if Azure Form Recognizer credentials are provided
    if args.formrecognizerservice == None:
        print(
            "Error: Azure Form Recognizer service is not provided. Please provide formrecognizerservice or use --localpdfparser for local pypdf parser.")
        exit(1)
    formrecognizer_creds = default_creds if args.formrecognizerkey == None and default_creds is not None else AzureKeyCredential(
        args.formrecognizerkey)

def blob_name_from_file_page(filename, page=0):
    if os.path.splitext(filename)[1].lower() == ".pdf":
        return os.path.splitext(os.path.basename(filename))[0] + f"-{page}" + ".pdf"
    else:
        return os.path.basename(filename)

def upload_blobs(filename):
    blob_service = BlobServiceClient(account_url=f"https://{args.storageaccount}.blob.core.windows.net",
                                     credential=storage_creds)
    blob_container = blob_service.get_container_client(args.container)
    if not blob_container.exists():
        blob_container.create_container()

    # if file is PDF split into pages and upload each page as a separate blob
    if os.path.splitext(filename)[1].lower() == ".pdf":
        reader = PdfReader(filename)
        pages = reader.pages
        for i in range(len(pages)):
            blob_name = blob_name_from_file_page(filename, i)
            if args.verbose: print(f"\tUploading blob for page {i} -> {blob_name}")
            f = io.BytesIO()
            writer = PdfWriter()
            writer.add_page(pages[i])
            writer.write(f)
            f.seek(0)
            blob_container.upload_blob(blob_name, f, overwrite=True)
    else:
        blob_name = blob_name_from_file_page(filename)
        with open(filename, "rb") as data:
            blob_container.upload_blob(blob_name, data, overwrite=True)

def remove_blobs(filename):
    if args.verbose: print(f"Removing blobs for '{filename or '<all>'}'")
    blob_service = BlobServiceClient(account_url=f"https://{args.storageaccount}.blob.core.windows.net",
                                     credential=storage_creds)
    blob_container = blob_service.get_container_client(args.container)
    if blob_container.exists():
        if filename == None:
            blobs = blob_container.list_blob_names()
        else:
            prefix = os.path.splitext(os.path.basename(filename))[0]
            blobs = filter(lambda b: re.match(f"{prefix}-\d+\.pdf", b), blob_container.list_blob_names(
                name_starts_with=os.path.splitext(os.path.basename(prefix))[0]))
        for b in blobs:
            if args.verbose: print(f"\tRemoving blob {b}")
            blob_container.delete_blob(b)

def table_to_html(table):
    table_html = "
"
    rows = [sorted([cell for cell in table.cells if cell.row_index == i], key=lambda cell: cell.column_index) for i in
            range(table.row_count)]
    for row_cells in rows:
        table_html += "
"
        for cell in row_cells:
            tag = "th" if (cell.kind == "columnHeader" or cell.kind == "rowHeader") else "td"
            cell_spans = ""
            if cell.column_span > 1: cell_spans += f" colSpan={cell.column_span}"
            if cell.row_span > 1: cell_spans += f" rowSpan={cell.row_span}"
            table_html += f"<{tag}{cell_spans}>{html.escape(cell.content)}</{tag}>"
        table_html += "
"
    table_html += "
"
    return table_html


from azure.identity import DefaultAzureCredential

from azure.search.documents import SearchClient

from azure.search.documents.models import QueryType

# Replace these with your own values, either in environment variables or directly here

AZURE_STORAGE_ACCOUNT = os.environ.get("AZURE_STORAGE_ACCOUNT") or "mystorageaccount"

AZURE_STORAGE_CONTAINER = os.environ.get("AZURE_STORAGE_CONTAINER") or "content"

AZURE_SEARCH_SERVICE = os.environ.get("AZURE_SEARCH_SERVICE") or "gptkb"

AZURE_SEARCH_INDEX = os.environ.get("AZURE_SEARCH_INDEX") or "gptkbindex"

AZURE_OPENAI_SERVICE = os.environ.get("AZURE_OPENAI_SERVICE") or "myopenai"

AZURE_OPENAI_GPT_DEPLOYMENT = os.environ.get("AZURE_OPENAI_GPT_DEPLOYMENT") or "davinci"

AZURE_OPENAI_CHATGPT_DEPLOYMENT = os.environ.get("AZURE_OPENAI_CHATGPT_DEPLOYMENT") or "chat"

KB_FIELDS_CONTENT = os.environ.get("KB_FIELDS_CONTENT") or "content"

KB_FIELDS_CATEGORY = os.environ.get("KB_FIELDS_CATEGORY") or "category"

KB_FIELDS_SOURCEPAGE = os.environ.get("KB_FIELDS_SOURCEPAGE") or "sourcepage"

# Use the current user identity to authenticate with Azure OpenAI, Cognitive Search and Blob Storage (no secrets needed,

# just use 'az login' locally, and managed identity when deployed on Azure). If you need to use keys, use separate AzureKeyCredential instances with the

# Comment these two lines out if using keys, set your API key in the OPENAI_API_KEY environment variable instead

# Set up clients for Cognitive Search and Storage

# ChatGPT uses a particular set of tokens to indicate turns in conversations

Assistant helps the company employees with their healthcare plan questions and employee handbook questions.

Answer ONLY with the facts listed in the list of sources below. If there isn't enough information below, say you don't know. Do not generate answers that don't use the sources below. If asking a clarifying question to the user would help, ask the question.

Each source has a name followed by colon and the actual information, always include the source name for each fact you use in the response. Use square brakets to reference the source, e.g. [info1.txt]. Don't combine sources, list each source separately, e.g. [info1.txt][info2.pdf].

summary_prompt_template = """Below is a summary of the conversation so far, and a new question asked by the user that needs to be answered by searching in a knowledge base. Generate a search query based on the conversation and the new question. Source names are not good search terms to include in the search query.

# Execute this cell multiple times updating user_input to accumulate chat history

user_input = "Does my plan cover annual eye exams?"

# Exclude category, to simulate scenarios where there's a set of docs you can't see

# Alternatively simply use search_client.search(q, top=3) if not using semantic search

filter = "category ne '{}'".format(exclude_category.replace("'", "''")) if exclude_category else None

results = [doc[KB_FIELDS_SOURCEPAGE] + ": " + doc[KB_FIELDS_CONTENT].replace("\n", "").replace("\r", "") for doc in r]

prompt = prompt_prefix.format(sources=content) + prompt_history + user_input + turn_suffix

prompt_history += user_input + turn_suffix + completion.choices[0].text + "\n<|im_end|>" + turn_prefix

history.append("assistant: " + completion.choices[0].text)


import json
import os
import tempfile
import uuid
import logging

from azure.core.credentials import AzureKeyCredential
from azure.core.exceptions import ResourceNotFoundError
from azure.search.documents.indexes.aio import SearchIndexClient
from azure.search.documents.aio import SearchClient
from azure.search.documents.indexes.models import (
    ComplexField,
    CorsOptions,
    SearchIndex,
    ScoringProfile,
    SearchFieldDataType,
    SimpleField,
    SearchableField
)

temp_path= "/tmp/"
logger = logging.getLogger('azure.search')
logger.setLevel(logging.DEBUG)

class AzureSearchClient(object):

    def __init__(self):
        self._service_endpoint = os.getenv("SEARCH_ENDPOINT")
        key = os.getenv("SEARCH_API_KEY")
        self._credential = AzureKeyCredential(key)
        self._admin_client = SearchIndexClient(endpoint=self._service_endpoint, credential=self._credential, logging_enable=True)
        self._index_client_map = {}

    async def close(self):
        await self._admin_client.close()
        for index_name, search_client in self._index_client_map.items( ):
            await search_client.close()

    async def create_index(self, index_name:str):
        print(f"Creating index: {index_name}")
        name = index_name
        fields = [
            SimpleField(name="id", type=SearchFieldDataType.String, filterable=True, sortable=True, key=True),
            SimpleField(name="date", type=SearchFieldDataType.DateTimeOffset, filterable=True),
            SimpleField(name="location", type=SearchFieldDataType.String, filterable=True),
            SimpleField(name="profile_name", type=SearchFieldDataType.String),
            SimpleField(name="rating", type=SearchFieldDataType.String, filterable=True),
            SearchableField(name="review_title", type=SearchFieldDataType.String, analyzer_name='en.microsoft'),
            SimpleField(name="stars", type=SearchFieldDataType.String, filterable=True),
            SimpleField(name="sentiment_aspects_NL", type="Collection(Edm.String)", filterable=True),
            SimpleField(name="aspects", type="Collection(Edm.String)", filterable=True),
            SimpleField(name="aspects_and_sentiments", type="Collection(Edm.String)", filterable=True),
            SearchableField(name="review_text", type=SearchFieldDataType.String, analyzer_name='en.microsoft')
        ]
        cors_options = CorsOptions(allowed_origins=["*"], max_age_in_seconds=60)
        scoring_profiles = []
        index = SearchIndex(
            name=name,
            fields=fields,
            scoring_profiles=scoring_profiles,
            cors_options=cors_options)
        try:
            result = await self._admin_client.create_index(index)
            print(f"Created index: {result}")
            print(f"Created index fields: {result.fields}")
        except Exception as e:
            print (e)
            raise Exception(f"Failed to create index {index_name}")

        # Preload search client for each index
        self._index_client_map[index_name] = self._admin_client.get_search_client(index_name)
        return True

    async def exists_index(self, index_name:str):
        try:
            await self._admin_client.get_index(index_name)
            self._index_client_map[index_name] = self._admin_client.get_search_client(index_name)
            return True
        except ResourceNotFoundError as ex:
            return False

    async def list_index(self):
        return [x for x in self._index_client_map.keys()]

    async def delete_index(self, index_name:str):
        try:
            result = await self._admin_client.delete_index(index_name)
            self._index_client_map.pop(index_name)
            return True
        except:
            return False

    async def upload_documents(self, index_name:str, documents:list[dict]):
        if await self.exists_index(index_name):
            index_client = self._index_client_map[index_name]
            print (index_client)
            try:
                result = await index_client.upload_documents(documents)
            except Exception as e:
                print (e)
            print (result)
            return result[0].succeeded
        return False

    async def search_documents(self, index_name:str, parameters:dict, top = 3):
        if await self.exists_index(index_name=index_name):
            index_client = self._index_client_map[index_name]
            print (f"Index found: {index_name}")
            try:
                print("Search client search documents parameters {}".format(parameters))
                documents = []
                async for document in await index_client.search(search_text= parameters['search_text'], filter = parameters['filter_text'], top = top, include_total_count=True):
                    documents.append(document)
                print("Num documents fetched from Az Search: {}".format(len(documents)))
                return documents
            except Exception as e:
                print (e)
            return None


ow
- The following figure shows an overview of the solution

![](./docs/AOAI_VA_Overview.png "Overview of the VA workflow")

## Workflow components

* Topic Classifier: This is responsible to determine the topic of the user input. The topic not only includes the business-related concepts but extends to include identifying if the input is a chitchat and if it is irrelevant to the scope of this virtual assistant. It also can be tuned to determine state like when the conversation is coming to an end. The topic classifier is a few shot model that uses the text-davinci-003 GPT base model. This
model is effective for classification tasks and is very powerful. It might be worth noting that a smaller, and cheaper model could be developed for a specific enterprise using their own data to perform the topic classification task.

* Text Generation: This component leverages ChatGPT to both understand the context and generate personalized responses to the current conversation. The exact prompt and how to form it will be discussed later in this file. The ChatGPT, using gpt-3.5-turbo as the base model, is used in this demonstration. It leverages the ability to both extract relevant information from the prompt and respond to a multi-turn conversation appropriately. The full context of the user and bot interaction, relevant to the current topic, is used in the prompt in addition to the relevant information surfaced from a knowledge base for the specific topic.

In addition, a dynamic tone modulator could be added, to help guide the tone and dynamic responses (through changing the ChatGPT parameters) in response to the user input and the current state. This would be injected into the prompt to help tune the output accordingly. In the current prompt used in this demo the tone is preset in the prompt to be "friendly".

# Run Time Architecture

The following figure illustrates the architecture of the runtime solution that enables the virtual assistant solution.

![](./docs/VA%20Runtime%20Arch.png "Archeticture of VA Runtime solution")

## Architecture Details
The runtime architecture of the chatbot includes a few components that enable a chatbot implementation. The core components includes:

* Orchestration: This module interfaces with the bot layer to exchange the user input as well as the response. The orchestrator retains the conversation state that is assessed during the dialog. This includes the topic of the current input as well as the full
transcript and session history. The transcript includes the response by ChatGPT in addition to the user input. This history is used in the generation of the prompt with the relevant context. The orchestrator also has access to the appropriate knowledge base to include in the prompt. With the relevant data, the orchestrator calls the model manager module to perform the API call to GPT or to ChatGPT.

* Model Manager: This module is responsible for forming the prompt, with the information elements passed to the model manager as well as the context and prompting information. The model manager combines the different pieces to provide a personalized prompt with the right context and transcript. The model manager also understands how to parse the ChatGPT response, and extract only one response that it would share with the orchestration layer to send to the user and retain in the session memory store.

* Cognition: This handles the GPT or ChatGPT API calls based on the prompt passed from the model manager. It returns the response to the manager to parse.

* Bot layer: This is a shell layer for managing the user interaction. For the purpose of this demonstration, this bot shell is based on the Bot Framework.

* User Interaction layer: This is a web app that hosts web chat control that is connected to the bot layer via web chat channel. User uses the web chat control to interact with the bot. The web app also has minimal controls to aid in this demo.

* Data Store: The user data and the knowledge base is stored in a Azure Cosmos DB, and is retrieved based on the user selected and/or the topic of the user input.

* Session history: The session history, for each user for a current session (identified by a consistent conversationID) is retained in an Azure Cosmos DB. The user input, in addition to the topic identified, as well as the bot response are all retained in the session history. This history is used in multiturn conversation to provide ChatGPT with context if needed.

# Exampl


import { AsyncCaller, AsyncCallerParams } from "./utils/async_caller.js";

/**
 * The parameters required to initialize an instance of the Embeddings
 * class.
 */
export type EmbeddingsParams = AsyncCallerParams;

export interface EmbeddingsInterface {
  /**
   * An abstract method that takes an array of documents as input and
   * returns a promise that resolves to an array of vectors for each
   * document.
   * @param documents An array of documents to be embedded.
   * @returns A promise that resolves to an array of vectors for each document.
   */
  embedDocuments(documents: string[]): Promise<number[][]>;

  /**
   * An abstract method that takes a single document as input and returns a
   * promise that resolves to a vector for the query document.
   * @param document A single document to be embedded.
   * @returns A promise that resolves to a vector for the query document.
   */
  embedQuery(document: string): Promise<number[]>;
}

/**
 * An abstract class that provides methods for embedding documents and
 * queries using LangChain.
 */
export abstract class Embeddings implements EmbeddingsInterface {
  /**
   * The async caller should be used by subclasses to make any async calls,
   * which will thus benefit from the concurrency and retry logic.
   */
  caller: AsyncCaller;

  constructor(params: EmbeddingsParams) {
    this.caller = new AsyncCaller(params ?? {});
  }

  /**
   * An abstract method that takes an array of documents as input and
   * returns a promise that resolves to an array of vectors for each
   * document.
   * @param documents An array of documents to be embedded.
   * @returns A promise that resolves to an array of vectors for each document.
   */
  abstract embedDocuments(documents: string[]): Promise<number[][]>;

  /**
   * An abstract method that takes a single document as input and returns a
   * promise that resolves to a vector for the query document.
   * @param document A single document to be embedded.
   * @returns A promise that resolves to a vector for the query document.
   */
  abstract embedQuery(document: string): Promise<number[]>;
}


/**
 * Type alias for a record where the keys are strings and the values can
 * be any type. This is used to represent the input values for a Chain.
 */
// eslint-disable-next-line @typescript-eslint/no-explicit-any
export type InputValues = Record<string, any>;

/**
 * Type alias for a record where the keys are strings and the values can
 * be any type. This is used to represent the output values from a Chain.
 */
// eslint-disable-next-line @typescript-eslint/no-explicit-any
export type OutputValues = Record<string, any>;

/**
 * Type alias for a record where the keys are strings and the values can
 * be any type. This is used to represent the memory variables in a Chain.
 */
// eslint-disable-next-line @typescript-eslint/no-explicit-any
export type MemoryVariables = Record<string, any>;

/**
 * Abstract base class for memory in LangChain's Chains. Memory refers to
 * the state in Chains. It can be used to store information about past
 * executions of a Chain and inject that information into the inputs of
 * future executions of the Chain.
 */
export abstract class BaseMemory {
  abstract get memoryKeys(): string[];

  /**
   * Abstract method that should take an object of input values and return a
   * Promise that resolves with an object of memory variables. The
   * implementation of this method should load the memory variables from the
   * provided input values.
   * @param values An object of input values.
   * @returns Promise that resolves with an object of memory variables.
   */
  abstract loadMemoryVariables(values: InputValues): Promise<MemoryVariables>;

  /**
   * Abstract method that should take two objects, one of input values and
   * one of output values, and return a Promise that resolves when the
   * context has been saved. The implementation of this method should save
   * the context based on the provided input and output values.
   * @param inputValues An object of input values.
   * @param outputValues An object of output values.
   * @returns Promise that resolves when the context has been saved.
   */
  abstract saveContext(
    inputValues: InputValues,
    outputValues: OutputValues
  ): Promise<void>;
}

const getValue = (values: InputValues | OutputValues, key?: string) => {
  if (key !== undefined) {
    return values[key];
  }
  const keys = Object.keys(values);
  if (keys.length === 1) {
    return values[keys[0]];
  }
};

/**
 * This function is used by memory classes to select the input value
 * to use for the memory. If there is only one input value, it is used.
 * If there are multiple input values, the inputKey must be specified.
 */
export const getInputValue = (inputValues: InputValues, inputKey?: string) => {
  const value = getValue(inputValues, inputKey);
  if (!value) {
    const keys = Object.keys(inputValues);
    throw new Error(
      `input values have ${keys.length} keys, you must specify an input key or pass only 1 key as input`
    );
  }
  return value;
};

/**
 * This function is used by memory classes to select the output value
 * to use for the memory. If there is only one output value, it is used.
 * If there are multiple output values, the outputKey must be specified.
 * If no outputKey is specified, an error is thrown.
 */
export const getOutputValue = (
  outputValues: OutputValues,
  outputKey?: string
) => {
  const value = getValue(outputValues, outputKey);
  if (!value) {
    const keys = Object.keys(outputValues);
    throw new Error(
      `output values have ${keys.length} keys, you must specify an output key or pass only 1 key as output`
    );
  }
  return value;
};

/**
 * Function used by memory classes to get the key of the prompt input,
 * excluding any keys that are memory variables or the "stop" key. If
 * there is not exactly one prompt input key, an error is thrown.
 */
export function getPromptInputKey(
  inputs: Record<string, unknown>,
  memoryVariables: string[]
): string {
  const promptInputKeys = Object.keys(inputs).filter(
    (key) => !memoryVariables.includes(key) && key !== "stop"
  );
  if (promptInputKeys.length !== 1) {
    throw new Error(
      `One input key expected, but got ${promptInputKeys.length}`
    );
  }
  return promptInputKeys[0];
}


import { test, expect } from "@jest/globals";
import { calculateMaxTokens, getModelContextSize } from "../base.js";

test("properly calculates correct max tokens", async () => {
  expect(
    await calculateMaxTokens({ prompt: "", modelName: "gpt-3.5-turbo-16k" })
  ).toBe(16384);
  expect(
    await calculateMaxTokens({
      prompt: "",
      modelName: "gpt-3.5-turbo-16k-0613",
    })
  ).toBe(16384);

  expect(
    await calculateMaxTokens({ prompt: "", modelName: "gpt-3.5-turbo" })
  ).toBe(4096);

  expect(await calculateMaxTokens({ prompt: "", modelName: "gpt-4" })).toBe(
    8192
  );
  expect(await calculateMaxTokens({ prompt: "", modelName: "gpt-4-32k" })).toBe(
    32768
  );
});

test("properly gets model context size", async () => {
  expect(await getModelContextSize("gpt-3.5-turbo-16k")).toBe(16384);
  expect(await getModelContextSize("gpt-3.5-turbo-16k-0613")).toBe(16384);
  expect(await getModelContextSize("gpt-3.5-turbo")).toBe(4096);
  expect(await getModelContextSize("gpt-4")).toBe(8192);
  expect(await getModelContextSize("gpt-4-32k")).toBe(32768);
});


export function mapStoredMessagesToChatMessages(
  messages: StoredMessage[]
): BaseMessage[] {
  return messages.map(mapStoredMessageToChatMessage);
}

/**
 * Transforms an array of `BaseMessage` instances into an array of
 * `StoredMessage` instances. It does this by calling the `toDict` method
 * on each `BaseMessage`, which returns a `StoredMessage`. This function
 * is used to prepare chat messages for storage.
 */
export function mapChatMessagesToStoredMessages(
  messages: BaseMessage[]
): StoredMessage[] {
  return messages.map((message) => message.toDict());
}

export function convertToChunk(message: BaseMessage) {
  const type = message._getType();
  if (type === "human") {
    // eslint-disable-next-line @typescript-eslint/no-use-before-define
    return new HumanMessageChunk({ ...message });
  } else if (type === "ai") {
    let aiChunkFields: AIMessageChunkFields = {
      ...message,
    };
    if ("tool_calls" in aiChunkFields) {
      aiChunkFields = {
        ...aiChunkFields,
        tool_call_chunks: aiChunkFields.tool_calls?.map((tc) => ({
          ...tc,
          type: "tool_call_chunk",
          index: undefined,
          args: JSON.stringify(tc.args),
        })),
      };
    }
    // eslint-disable-next-line @typescript-eslint/no-use-before-define
    return new AIMessageChunk({ ...aiChunkFields });
  } else if (type === "system") {
    // eslint-disable-next-line @typescript-eslint/no-use-before-define
    return new SystemMessageChunk({ ...message });
  } else if (type === "function") {
    // eslint-disable-next-line @typescript-eslint/no-use-before-define
    return new FunctionMessageChunk({ ...message });
    // eslint-disable-next-line @typescript-eslint/no-use-before-define
  } else if (ChatMessage.isInstance(message)) {
    // eslint-disable-next-line @typescript-eslint/no-use-before-define
    return new ChatMessageChunk({ ...message });
  } else {
    throw new Error("Unknown message type.");
  }
}


import {
  BaseMessage,
  BaseMessageChunk,
  mergeContent,
  _mergeDicts,
  type MessageType,
} from "./base.js";

/**
 * Represents a system message in a conversation.
 */
export class SystemMessage extends BaseMessage {
  static lc_name() {
    return "SystemMessage";
  }

  _getType(): MessageType {
    return "system";
  }
}

/**
 * Represents a chunk of a system message, which can be concatenated with
 * other system message chunks.
 */
export class SystemMessageChunk extends BaseMessageChunk {
  static lc_name() {
    return "SystemMessageChunk";
  }

  _getType(): MessageType {
    return "system";
  }

  concat(chunk: SystemMessageChunk) {
    return new SystemMessageChunk({
      content: mergeContent(this.content, chunk.content),
      additional_kwargs: _mergeDicts(
        this.additional_kwargs,
        chunk.additional_kwargs
      ),
      response_metadata: _mergeDicts(
        this.response_metadata,
        chunk.response_metadata
      ),
      id: this.id ?? chunk.id,
    });
  }
}

export function isSystemMessage(x: BaseMessage): x is SystemMessage {
  return x._getType() === "system";
}

export function isSystemMessageChunk(
  x: BaseMessageChunk
): x is SystemMessageChunk {
  return x._getType() === "system";
}


import { it, describe, test, expect } from "@jest/globals";
import {
  filterMessages,
  mergeMessageRuns,
  trimMessages,
} from "../transformers.js";
import { AIMessage } from "../ai.js";
import { ChatMessage } from "../chat.js";
import { HumanMessage } from "../human.js";
import { SystemMessage } from "../system.js";
import { BaseMessage } from "../base.js";
import {
  getBufferString,
  mapChatMessagesToStoredMessages,
  mapStoredMessagesToChatMessages,
} from "../utils.js";

describe("filterMessage", () => {
  const getMessages = () => [
    new SystemMessage("you're a good assistant."),
    new HumanMessage({
      content: "what's your name",
      id: "foo",
      name: "example_user",
    }),
    new AIMessage({ content: "steve-o", id: "bar", name: "example_assistant" }),
    new HumanMessage({ content: "what's your favorite color", id: "baz" }),
    new AIMessage({ content: "silicon blue", id: "blah" }),
  ];

  it("works", () => {
    const messages = getMessages();
    const filteredMessages = filterMessages(messages, {
      includeNames: ["example_user", "example_assistant"],
      includeTypes: ["system"],
      excludeIds: ["bar"],
    });
    expect(filteredMessages).toEqual([
      new SystemMessage("you're a good assistant."),
      new HumanMessage({
        content: "what's your name",
        id: "foo",
        name: "example_user",
      }),
    ]);
  });

  it("can filter messages based on class types", () => {
    const messages = getMessages();

    const filteredMessages = filterMessages(messages, {
      includeTypes: [HumanMessage, AIMessage],
    });
    expect(filteredMessages).toHaveLength(4);
    expect(filteredMessages).toEqual([
      new HumanMessage({
        content: "what's your name",
        id: "foo",
        name: "example_user",
      }),
      new AIMessage({
        content: "steve-o",
        id: "bar",
        name: "example_assistant",
      }),
      new HumanMessage({ content: "what's your favorite color", id: "baz" }),
      new AIMessage({ content: "silicon blue", id: "blah" }),
    ]);
  });

  it("returns a runnable if no messages are passed", () => {
    const filteredMessagesRunnable = filterMessages();
    expect(filteredMessagesRunnable).toBeDefined();
    expect(filteredMessagesRunnable.lc_namespace).toEqual([
      "langchain_core",
      "runnables",
    ]);
    expect("func" in filteredMessagesRunnable).toBeTruthy();
    // `func` is protected, so we need to cast it to any to access it
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    expect(typeof (filteredMessagesRunnable as any).func).toBe("function");
  });
});

describe("mergeMessageRuns", () => {
  const getMessages = () => [
    new SystemMessage("you're a good assistant."),
    new HumanMessage({ content: "what's your favorite color", id: "foo" }),
    new HumanMessage({ content: "wait your favorite food", id: "bar" }),
    new AIMessage({
      content: "my favorite colo",
      tool_calls: [{ name: "blah_tool", args: { x: 2 }, id: "123" }],
      id: "baz",
    }),
    new AIMessage({
      content: [{ type: "text", text: "my favorite dish is lasagna" }],
      tool_calls: [{ name: "blah_tool", args: { x: -10 }, id: "456" }],
      id: "blur",
    }),
  ];

  it("works", () => {
    const messages = getMessages();

    const mergedMessages = mergeMessageRuns(messages);
    expect(mergedMessages).toHaveLength(3);
    expect(mergedMessages).toEqual([
      new SystemMessage("you're a good assistant."),
      new HumanMessage({
        content: "what's your favorite color\nwait your favorite food",
        id: "foo",
      }),
      new AIMessage({
        content: [
          { type: "text", text: "my favorite colo" },
          { type: "text", text: "my favorite dish is lasagna" },
        ],
        tool_calls: [
          { name: "blah_tool", args: { x: 2 }, id: "123", type: "tool_call" },
          { name: "blah_tool", args: { x: -10 }, id: "456", type: "tool_call" },
        ],
        id: "baz",
      }),
    ]);
  });

  it("returns a runnable if no messages are passed", () => {
    const mergedMessages = mergeMessageRuns();
    expect(mergedMessages).toBeDefined();
    expect(mergedMessages.lc_namespace).toEqual([
      "langchain_core",
      "runnables",
    ]);
    expect("func" in mergedMessages).toBeTruthy();
    // `func` is protected, so we need to cast it to any to access it
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    expect(typeof (mergedMessages as any).func).toBe("function");
  });
});


describe("trimMessages can trim", () => {
  const messagesAndTokenCounterFactory = () => {
    const messages = [
      new SystemMessage(
        "This is a 4 token text. The full message is 10 tokens."
      ),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "first",
      }),
      new AIMessage({
        content: [
          { type: "text", text: "This is the FIRST 4 token block." },
          { type: "text", text: "This is the SECOND 4 token block." },
        ],
        id: "second",
      }),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "third",
      }),
      new AIMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "fourth",
      }),
    ];

    const dummyTokenCounter = (messages: BaseMessage[]): number => {
      // treat each message like it adds 3 default tokens at the beginning
      // of the message and at the end of the message. 3 + 4 + 3 = 10 tokens
      // per message.

      const defaultContentLen = 4;
      const defaultMsgPrefixLen = 3;
      const defaultMsgSuffixLen = 3;

      let count = 0;
      for (const msg of messages) {
        if (typeof msg.content === "string") {
          count +=
            defaultMsgPrefixLen + defaultContentLen + defaultMsgSuffixLen;
        }
        if (Array.isArray(msg.content)) {
          count +=
            defaultMsgPrefixLen +
            msg.content.length * defaultContentLen +
            defaultMsgSuffixLen;
        }
      }
      return count;
    };

    return {
      messages,
      dummyTokenCounter,
    };
  };

  it("First 30 tokens, not allowing partial messages", async () => {
    const { messages, dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = await trimMessages(messages, {
      maxTokens: 30,
      tokenCounter: dummyTokenCounter,
      strategy: "first",
    });
    expect(trimmedMessages).toHaveLength(2);
    expect(trimmedMessages).toEqual([
      new SystemMessage(
        "This is a 4 token text. The full message is 10 tokens."
      ),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "first",
      }),
    ]);
  });

  it("First 30 tokens, allowing partial messages", async () => {
    const { messages, dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = await trimMessages(messages, {
      maxTokens: 30,
      tokenCounter: dummyTokenCounter,
      strategy: "first",
      allowPartial: true,
    });

    expect(trimmedMessages).toHaveLength(3);
    expect(trimmedMessages).toEqual([
      new SystemMessage(
        "This is a 4 token text. The full message is 10 tokens."
      ),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "first",
      }),
      new AIMessage({
        content: [{ type: "text", text: "This is the FIRST 4 token block." }],
        id: "second",
      }),
    ]);
  });

  it("First 30 tokens, allowing partial messages, have to end on HumanMessage", async () => {
    const { messages, dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = await trimMessages(messages, {
      maxTokens: 30,
      tokenCounter: dummyTokenCounter,
      strategy: "first",
      allowPartial: true,
      endOn: "human",
    });

    expect(trimmedMessages).toHaveLength(2);
    expect(trimmedMessages).toEqual([
      new SystemMessage(
        "This is a 4 token text. The full message is 10 tokens."
      ),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "first",
      }),
    ]);
  });

  it("Last 30 tokens, including system message, not allowing partial messages", async () => {
    const { messages, dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = await trimMessages(messages, {
      maxTokens: 30,
      includeSystem: true,
      tokenCounter: dummyTokenCounter,
      strategy: "last",
    });

    expect(trimmedMessages).toHaveLength(3);
    expect(trimmedMessages).toEqual([
      new SystemMessage(
        "This is a 4 token text. The full message is 10 tokens."
      ),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "third",
      }),
      new AIMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "fourth",
      }),
    ]);
  });

  it("Last 40 tokens, including system message, allowing partial messages", async () => {
    const { messages, dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = await trimMessages(messages, {
      maxTokens: 40,
      tokenCounter: dummyTokenCounter,
      strategy: "last",
      allowPartial: true,
      includeSystem: true,
    });

    expect(trimmedMessages).toHaveLength(4);
    expect(trimmedMessages).toEqual([
      new SystemMessage(
        "This is a 4 token text. The full message is 10 tokens."
      ),
      new AIMessage({
        content: [{ type: "text", text: "This is the FIRST 4 token block." }],
        id: "second",
      }),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "third",
      }),
      new AIMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "fourth",
      }),
    ]);
  });

  it("Last 30 tokens, including system message, allowing partial messages, end on HumanMessage", async () => {
    const { messages, dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = await trimMessages(messages, {
      maxTokens: 30,
      tokenCounter: dummyTokenCounter,
      strategy: "last",
      endOn: "human",
      includeSystem: true,
      allowPartial: true,
    });

    expect(trimmedMessages).toHaveLength(3);
    expect(trimmedMessages).toEqual([
      new SystemMessage(
        "This is a 4 token text. The full message is 10 tokens."
      ),
      new AIMessage({
        content: [{ type: "text", text: "This is the FIRST 4 token block." }],
        id: "second",
      }),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "third",
      }),
    ]);
  });

  it("Last 40 tokens, including system message, allowing partial messages, start on HumanMessage", async () => {
    const { messages, dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = await trimMessages(messages, {
      maxTokens: 40,
      tokenCounter: dummyTokenCounter,
      strategy: "last",
      includeSystem: true,
      allowPartial: true,
      startOn: "human",
    });

    expect(trimmedMessages).toHaveLength(3);
    console.log(trimmedMessages);
    expect(trimmedMessages).toEqual([
      new SystemMessage(
        "This is a 4 token text. The full message is 10 tokens."
      ),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "third",
      }),
      new AIMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "fourth",
      }),
    ]);
  });

  it("can filter (startOn) with message classes", async () => {
    const { messages, dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = await trimMessages(messages, {
      maxTokens: 40,
      tokenCounter: dummyTokenCounter,
      startOn: [HumanMessage],
    });
    expect(trimmedMessages).toHaveLength(2);
    expect(trimmedMessages).toEqual([
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "third",
      }),
      new AIMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "fourth",
      }),
    ]);
  });


it("can filter (endOn) with message classes", async () => {
    const { messages, dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = await trimMessages(messages, {
      maxTokens: 40,
      tokenCounter: dummyTokenCounter,
      endOn: [HumanMessage],
    });
    console.log(trimmedMessages);
    expect(trimmedMessages).toHaveLength(3);
    expect(trimmedMessages).toEqual([
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "first",
      }),
      new AIMessage({
        content: [
          { type: "text", text: "This is the FIRST 4 token block." },
          { type: "text", text: "This is the SECOND 4 token block." },
        ],
        id: "second",
      }),
      new HumanMessage({
        content: "This is a 4 token text. The full message is 10 tokens.",
        id: "third",
      }),
    ]);
  });

  it("can return a runnable if empty array is passed", () => {
    const { dummyTokenCounter } = messagesAndTokenCounterFactory();
    const trimmedMessages = trimMessages({
      maxTokens: 40,
      tokenCounter: dummyTokenCounter,
    });
    expect(trimmedMessages).toBeDefined();
    expect(trimmedMessages.lc_namespace).toEqual([
      "langchain_core",
      "runnables",
    ]);
    expect("func" in trimmedMessages).toBeTruthy();
    // `func` is protected, so we need to cast it to any to access it
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    expect(typeof (trimmedMessages as any).func).toBe("function");
  });
});

test("getBufferString can handle complex messages", () => {
  const messageArr1 = [new HumanMessage("Hello there!")];
  const messageArr2 = [
    new AIMessage({
      content: [
        {
          type: "text",
          text: "Hello there!",
        },
      ],
    }),
  ];
  const messageArr3 = [
    new HumanMessage({
      content: [
        {
          type: "image_url",
          image_url: {
            url: "https://example.com/image.jpg",
          },
        },
        {
          type: "image_url",
          image_url: "https://example.com/image.jpg",
        },
      ],
    }),
  ];

  const bufferString1 = getBufferString(messageArr1);
  expect(bufferString1).toBe("Human: Hello there!");

  const bufferString2 = getBufferString(messageArr2);
  expect(bufferString2).toBe(
    `AI: ${JSON.stringify(
      [
        {
          type: "text",
          text: "Hello there!",
        },
      ],
      null,
      2
    )}`
  );

  const bufferString3 = getBufferString(messageArr3);
  expect(bufferString3).toBe(
    `Human: ${JSON.stringify(
      [
        {
          type: "image_url",
          image_url: {
            url: "https://example.com/image.jpg",
          },
        },
        {
          type: "image_url",
          image_url: "https://example.com/image.jpg",
        },
      ],
      null,
      2
    )}`
  );
});

describe("chat message conversions", () => {
  it("can convert a chat message to a stored message and back", () => {
    const originalMessages = [
      new ChatMessage("I'm a generic message!", "human"),
      new HumanMessage("I'm a human message!"),
    ];

    const storedMessages = mapChatMessagesToStoredMessages(originalMessages);

    const convertedBackMessages =
      mapStoredMessagesToChatMessages(storedMessages);

    expect(convertedBackMessages).toEqual(originalMessages);
  });
});


// Default generic "any" values are for backwards compatibility.
// Replace with "string" when we are comfortable with a breaking change.

import { BaseStringPromptTemplate } from "./string.js";
import type {
  BasePromptTemplateInput,
  TypedPromptInputValues,
} from "./base.js";
import {
  checkValidTemplate,
  parseTemplate,
  renderTemplate,
  type TemplateFormat,
} from "./template.js";
import type { SerializedPromptTemplate } from "./serde.js";
import type { InputValues, PartialValues } from "../utils/types/index.js";
import { MessageContent, MessageContentComplex } from "../messages/index.js";

/**
 * Inputs to create a {@link PromptTemplate}
 * @augments BasePromptTemplateInput
 */
export interface PromptTemplateInput<
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  RunInput extends InputValues = any,
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  PartialVariableName extends string = any,
  Format extends TemplateFormat = TemplateFormat
> extends BasePromptTemplateInput<RunInput, PartialVariableName> {
  /**
   * The prompt template
   */
  template: MessageContent;

  /**
   * The format of the prompt template. Options are "f-string" and "mustache"
   */
  templateFormat?: Format;

  /**
   * Whether or not to try validating the template on initialization
   *
   * @defaultValue `true`
   */
  validateTemplate?: boolean;

  /**
   * Additional fields which should be included inside
   * the message content array if using a complex message
   * content.
   */
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  additionalContentFields?: MessageContentComplex;
}

type NonAlphanumeric =
  | " "
  | "\t"
  | "\n"
  | "\r"
  | '"'
  | "'"
  | "{"
  | "["
  | "("
  | "`"
  | ":"
  | ";";

/**
 * Recursive type to extract template parameters from a string.
 * @template T - The input string.
 * @template Result - The resulting array of extracted template parameters.
 */
type ExtractTemplateParamsRecursive<
  T extends string,
  Result extends string[] = []
> = T extends `${string}{${infer Param}}${infer Rest}`
  ? Param extends `${NonAlphanumeric}${string}`
    ? ExtractTemplateParamsRecursive<Rest, Result> // for non-template variables that look like template variables e.g. see https://github.com/langchain-ai/langchainjs/blob/main/langchain/src/chains/query_constructor/prompt.ts
    : ExtractTemplateParamsRecursive<Rest, [...Result, Param]>
  : Result;

export type ParamsFromFString<T extends string> = {
  [Key in
    | ExtractTemplateParamsRecursive<T>[number]
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    | (string & Record<never, never>)]: any;
};

export type ExtractedFStringParams<
  T extends string,
  // eslint-disable-next-line @typescript-eslint/ban-types
  RunInput extends InputValues = Symbol
  // eslint-disable-next-line @typescript-eslint/ban-types
> = RunInput extends Symbol ? ParamsFromFString<T> : RunInput;

/**
 * Schema to represent a basic prompt for an LLM.
 * @augments BasePromptTemplate
 * @augments PromptTemplateInput
 *
 * @example
 * ```ts
 * import { PromptTemplate } from "langchain/prompts";
 *
 * const prompt = new PromptTemplate({
 *   inputVariables: ["foo"],
 *   template: "Say {foo}",
 * });
 * ```
 */


export class PromptTemplate<
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    RunInput extends InputValues = any,
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    PartialVariableName extends string = any
  >
  extends BaseStringPromptTemplate<RunInput, PartialVariableName>
  implements PromptTemplateInput<RunInput, PartialVariableName>
{
  static lc_name() {
    return "PromptTemplate";
  }

  template: MessageContent;

  templateFormat: TemplateFormat = "f-string";

  validateTemplate = true;

  /**
   * Additional fields which should be included inside
   * the message content array if using a complex message
   * content.
   */
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  additionalContentFields?: MessageContentComplex;

  constructor(input: PromptTemplateInput<RunInput, PartialVariableName>) {
    super(input);
    // If input is mustache and validateTemplate is not defined, set it to false
    if (
      input.templateFormat === "mustache" &&
      input.validateTemplate === undefined
    ) {
      this.validateTemplate = false;
    }
    Object.assign(this, input);

    if (this.validateTemplate) {
      if (this.templateFormat === "mustache") {
        throw new Error("Mustache templates cannot be validated.");
      }
      let totalInputVariables: string[] = this.inputVariables;
      if (this.partialVariables) {
        totalInputVariables = totalInputVariables.concat(
          Object.keys(this.partialVariables)
        );
      }
      checkValidTemplate(
        this.template,
        this.templateFormat,
        totalInputVariables
      );
    }
  }

  _getPromptType(): "prompt" {
    return "prompt";
  }

  /**
   * Formats the prompt template with the provided values.
   * @param values The values to be used to format the prompt template.
   * @returns A promise that resolves to a string which is the formatted prompt.
   */
  async format(values: TypedPromptInputValues<RunInput>): Promise<string> {
    const allValues = await this.mergePartialAndUserVariables(values);
    return renderTemplate(
      this.template as string,
      this.templateFormat,
      allValues
    );
  }

  /**
   * Take examples in list format with prefix and suffix to create a prompt.
   *
   * Intended to be used a a way to dynamically create a prompt from examples.
   *
   * @param examples - List of examples to use in the prompt.
   * @param suffix - String to go after the list of examples. Should generally set up the user's input.
   * @param inputVariables - A list of variable names the final prompt template will expect
   * @param exampleSeparator - The separator to use in between examples
   * @param prefix - String that should go before any examples. Generally includes examples.
   *
   * @returns The final prompt template generated.
   */
  static fromExamples(
    examples: string[],
    suffix: string,
    inputVariables: string[],
    exampleSeparator = "\n\n",
    prefix = ""
  ) {
    const template = [prefix, ...examples, suffix].join(exampleSeparator);
    return new PromptTemplate({
      inputVariables,
      template,
    });
  }

  /**
   * Load prompt template from a template f-string
   */
  static fromTemplate<
    // eslint-disable-next-line @typescript-eslint/ban-types
    RunInput extends InputValues = Symbol,
    T extends string = string
  >(
    template: T,
    options?: Omit<
      PromptTemplateInput<RunInput, string, "f-string">,
      "template" | "inputVariables"
    >
  ): PromptTemplate<ExtractedFStringParams<T, RunInput>>;

  static fromTemplate<
    // eslint-disable-next-line @typescript-eslint/ban-types
    RunInput extends InputValues = Symbol,
    T extends string = string
  >(
    template: T,
    options?: Omit<
      PromptTemplateInput<RunInput, string>,
      "template" | "inputVariables"
    >
  ): PromptTemplate<ExtractedFStringParams<T, RunInput>>;

  static fromTemplate<
    // eslint-disable-next-line @typescript-eslint/ban-types
    RunInput extends InputValues = Symbol,
    T extends string = string
  >(
    template: T,
    options?: Omit<
      PromptTemplateInput<RunInput, string, "mustache">,
      "template" | "inputVariables"
    >
  ): PromptTemplate<InputValues>;

  static fromTemplate<
    // eslint-disable-next-line @typescript-eslint/ban-types
    RunInput extends InputValues = Symbol,
    T extends string = string
  >(
    template: T,
    options?: Omit<
      PromptTemplateInput<RunInput, string, TemplateFormat>,
      "template" | "inputVariables"
    >
  ): PromptTemplate<ExtractedFStringParams<T, RunInput> | InputValues> {
    const { templateFormat = "f-string", ...rest } = options ?? {};
    const names = new Set<string>();
    parseTemplate(template, templateFormat).forEach((node) => {
      if (node.type === "variable") {
        names.add(node.name);
      }
    });

    return new PromptTemplate({
      // Rely on extracted types
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      inputVariables: [...names] as any[],
      templateFormat,
      template,
      ...rest,
    });
  }

  /**
   * Partially applies values to the prompt template.
   * @param values The values to be partially applied to the prompt template.
   * @returns A new instance of PromptTemplate with the partially applied values.
   */
  async partial<NewPartialVariableName extends string>(
    values: PartialValues<NewPartialVariableName>
  ) {
    const newInputVariables = this.inputVariables.filter(
      (iv) => !(iv in values)
    ) as Exclude<Extract<keyof RunInput, string>, NewPartialVariableName>[];
    const newPartialVariables = {
      ...(this.partialVariables ?? {}),
      ...values,
    } as PartialValues<PartialVariableName | NewPartialVariableName>;
    const promptDict = {
      ...this,
      inputVariables: newInputVariables,
      partialVariables: newPartialVariables,
    };
    return new PromptTemplate<
      InputValues<
        Exclude<Extract<keyof RunInput, string>, NewPartialVariableName>
      >
    >(promptDict);
  }

  serialize(): SerializedPromptTemplate {
    if (this.outputParser !== undefined) {
      throw new Error(
        "Cannot serialize a prompt template with an output parser"
      );
    }
    return {
      _type: this._getPromptType(),
      input_variables: this.inputVariables,
      template: this.template,
      template_format: this.templateFormat,
    };
  }

  static async deserialize(
    data: SerializedPromptTemplate
  ): Promise<PromptTemplate> {
    if (!data.template) {
      throw new Error("Prompt template must have a template");
    }
    const res = new PromptTemplate({
      inputVariables: data.input_variables,
      template: data.template,
      templateFormat: data.template_format,
    });
    return res;
  }

  // TODO(from file)
}


export class ChatPromptTemplate<
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    RunInput extends InputValues = any,
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    PartialVariableName extends string = any
  >
  extends BaseChatPromptTemplate<RunInput, PartialVariableName>
  implements ChatPromptTemplateInput<RunInput, PartialVariableName>
{
  static lc_name() {
    return "ChatPromptTemplate";
  }

  get lc_aliases(): Record<string, string> {
    return {
      promptMessages: "messages",
    };
  }

  promptMessages: Array<BaseMessagePromptTemplate | BaseMessage>;

  validateTemplate = true;

  templateFormat: TemplateFormat = "f-string";

  constructor(input: ChatPromptTemplateInput<RunInput, PartialVariableName>) {
    super(input);
    // If input is mustache and validateTemplate is not defined, set it to false
    if (
      input.templateFormat === "mustache" &&
      input.validateTemplate === undefined
    ) {
      this.validateTemplate = false;
    }
    Object.assign(this, input);

    if (this.validateTemplate) {
      const inputVariablesMessages = new Set<string>();
      for (const promptMessage of this.promptMessages) {
        // eslint-disable-next-line no-instanceof/no-instanceof
        if (promptMessage instanceof BaseMessage) continue;
        for (const inputVariable of promptMessage.inputVariables) {
          inputVariablesMessages.add(inputVariable);
        }
      }

      const totalInputVariables = this.inputVariables as string[];
      const inputVariablesInstance = new Set(
        this.partialVariables
          ? totalInputVariables.concat(Object.keys(this.partialVariables))
          : totalInputVariables
      );
      const difference = new Set(
        [...inputVariablesInstance].filter(
          (x) => !inputVariablesMessages.has(x)
        )
      );
      if (difference.size > 0) {
        throw new Error(
          `Input variables \`${[
            ...difference,
          ]}\` are not used in any of the prompt messages.`
        );
      }
      const otherDifference = new Set(
        [...inputVariablesMessages].filter(
          (x) => !inputVariablesInstance.has(x)
        )
      );
      if (otherDifference.size > 0) {
        throw new Error(
          `Input variables \`${[
            ...otherDifference,
          ]}\` are used in prompt messages but not in the prompt template.`
        );
      }
    }
  }

  _getPromptType(): "chat" {
    return "chat";
  }

  private async _parseImagePrompts(
    message: BaseMessage,
    inputValues: InputValues<
      PartialVariableName | Extract<keyof RunInput, string>
    >
  ): Promise<BaseMessage> {
    if (typeof message.content === "string") {
      return message;
    }
    const formattedMessageContent = await Promise.all(
      message.content.map(async (item) => {
        if (item.type !== "image_url") {
          return item;
        }

        let imageUrl = "";
        if (typeof item.image_url === "string") {
          imageUrl = item.image_url;
        } else {
          imageUrl = item.image_url.url;
        }

        const promptTemplatePlaceholder = PromptTemplate.fromTemplate(
          imageUrl,
          {
            templateFormat: this.templateFormat,
          }
        );
        const formattedUrl = await promptTemplatePlaceholder.format(
          inputValues
        );

        if (typeof item.image_url !== "string" && "url" in item.image_url) {
          // eslint-disable-next-line no-param-reassign
          item.image_url.url = formattedUrl;
        } else {
          // eslint-disable-next-line no-param-reassign
          item.image_url = formattedUrl;
        }
        return item;
      })
    );
    // eslint-disable-next-line no-param-reassign
    message.content = formattedMessageContent;
    return message;
  }

  async formatMessages(
    values: TypedPromptInputValues<RunInput>
  ): Promise<BaseMessage[]> {
    const allValues = await this.mergePartialAndUserVariables(values);
    let resultMessages: BaseMessage[] = [];

    for (const promptMessage of this.promptMessages) {
      // eslint-disable-next-line no-instanceof/no-instanceof
      if (promptMessage instanceof BaseMessage) {
        resultMessages.push(
          await this._parseImagePrompts(promptMessage, allValues)
        );
      } else {
        const inputValues = promptMessage.inputVariables.reduce(
          (acc, inputVariable) => {
            if (
              !(inputVariable in allValues) &&
              !(isMessagesPlaceholder(promptMessage) && promptMessage.optional)
            ) {
              const error = addLangChainErrorFields(
                new Error(
                  `Missing value for input variable \`${inputVariable.toString()}\``
                ),
                "INVALID_PROMPT_INPUT"
              );
              throw error;
            }
            acc[inputVariable] = allValues[inputVariable];
            return acc;
          },
          {} as InputValues
        );
        const message = await promptMessage.formatMessages(inputValues);
        resultMessages = resultMessages.concat(message);
      }
    }
    return resultMessages;
  }

  async partial<NewPartialVariableName extends string>(
    values: PartialValues<NewPartialVariableName>
  ) {
    // This is implemented in a way it doesn't require making
    // BaseMessagePromptTemplate aware of .partial()
    const newInputVariables = this.inputVariables.filter(
      (iv) => !(iv in values)
    ) as Exclude<Extract<keyof RunInput, string>, NewPartialVariableName>[];
    const newPartialVariables = {
      ...(this.partialVariables ?? {}),
      ...values,
    } as PartialValues<PartialVariableName | NewPartialVariableName>;
    const promptDict = {
      ...this,
      inputVariables: newInputVariables,
      partialVariables: newPartialVariables,
    };
    return new ChatPromptTemplate<
      InputValues<
        Exclude<Extract<keyof RunInput, string>, NewPartialVariableName>
      >
    >(promptDict);
  }

  /**
   * Load prompt template from a template f-string
   */
  static fromTemplate<
    // eslint-disable-next-line @typescript-eslint/ban-types
    RunInput extends InputValues = Symbol,
    T extends string = string
  >(
    template: T,
    options?: Omit<
      PromptTemplateInput<RunInput, string, "f-string">,
      "template" | "inputVariables"
    >
  ): ChatPromptTemplate<ExtractedFStringParams<T, RunInput>>;

  static fromTemplate<
    // eslint-disable-next-line @typescript-eslint/ban-types
    RunInput extends InputValues = Symbol,
    T extends string = string
  >(
    template: T,
    options?: Omit<
      PromptTemplateInput<RunInput, string>,
      "template" | "inputVariables"
    >
  ): ChatPromptTemplate<ExtractedFStringParams<T, RunInput>>;

  static fromTemplate<
    // eslint-disable-next-line @typescript-eslint/ban-types
    RunInput extends InputValues = Symbol,
    T extends string = string
  >(
    template: T,
    options?: Omit<
      PromptTemplateInput<RunInput, string, "mustache">,
      "template" | "inputVariables"
    >
  ): ChatPromptTemplate<InputValues>;

  static fromTemplate<
    // eslint-disable-next-line @typescript-eslint/ban-types
    RunInput extends InputValues = Symbol,
    T extends string = string
  >(
    template: T,
    options?: Omit<
      PromptTemplateInput<RunInput, string, TemplateFormat>,
      "template" | "inputVariables"
    >
  ): ChatPromptTemplate<ExtractedFStringParams<T, RunInput> | InputValues> {
    const prompt = PromptTemplate.fromTemplate(template, options);
    const humanTemplate = new HumanMessagePromptTemplate({ prompt });
    return this.fromMessages<
      // eslint-disable-next-line @typescript-eslint/ban-types
      RunInput extends Symbol ? ParamsFromFString<T> : RunInput
    >([humanTemplate]);
  }

  /**
   * Create a chat model-specific prompt from individual chat messages
   * or message-like tuples.
   * @param promptMessages Messages to be passed to the chat model
   * @returns A new ChatPromptTemplate
   */


import { BaseStringPromptTemplate } from "./string.js";
import type {
  BasePromptTemplateInput,
  TypedPromptInputValues,
  Example,
} from "./base.js";
import type { BaseExampleSelector } from "../example_selectors/base.js";
import {
  type TemplateFormat,
  checkValidTemplate,
  renderTemplate,
} from "./template.js";
import { PromptTemplate } from "./prompt.js";
import type { SerializedFewShotTemplate } from "./serde.js";
import type { InputValues, PartialValues } from "../utils/types/index.js";
import type { BaseMessage } from "../messages/index.js";
import {
  BaseChatPromptTemplate,
  type BaseMessagePromptTemplate,
} from "./chat.js";

export interface FewShotPromptTemplateInput
  extends BasePromptTemplateInput<InputValues> {
  /**
   * Examples to format into the prompt. Exactly one of this or
   * {@link exampleSelector} must be
   * provided.
   */
  examples?: Example[];

  /**
   * An {@link BaseExampleSelector} Examples to format into the prompt. Exactly one of this or
   * {@link examples} must be
   * provided.
   */
  exampleSelector?: BaseExampleSelector;

  /**
   * An {@link PromptTemplate} used to format a single example.
   */
  examplePrompt: PromptTemplate;

  /**
   * String separator used to join the prefix, the examples, and suffix.
   */
  exampleSeparator?: string;

  /**
   * A prompt template string to put before the examples.
   *
   * @defaultValue `""`
   */
  prefix?: string;

  /**
   * A prompt template string to put after the examples.
   */
  suffix?: string;

  /**
   * The format of the prompt template. Options are: 'f-string'
   */
  templateFormat?: TemplateFormat;

  /**
   * Whether or not to try validating the template on initialization.
   */
  validateTemplate?: boolean;
}

/**
 * Prompt template that contains few-shot examples.
 * @augments BasePromptTemplate
 * @augments FewShotPromptTemplateInput
 * @example
 * ```typescript
 * const examplePrompt = PromptTemplate.fromTemplate(
 *   "Input: {input}\nOutput: {output}",
 * );
 *
 * const exampleSelector = await SemanticSimilarityExampleSelector.fromExamples(
 *   [
 *     { input: "happy", output: "sad" },
 *     { input: "tall", output: "short" },
 *     { input: "energetic", output: "lethargic" },
 *     { input: "sunny", output: "gloomy" },
 *     { input: "windy", output: "calm" },
 *   ],
 *   new OpenAIEmbeddings(),
 *   HNSWLib,
 *   { k: 1 },
 * );
 *
 * const dynamicPrompt = new FewShotPromptTemplate({
 *   exampleSelector,
 *   examplePrompt,
 *   prefix: "Give the antonym of every input",
 *   suffix: "Input: {adjective}\nOutput:",
 *   inputVariables: ["adjective"],
 * });
 *
 * // Format the dynamic prompt with the input 'rainy'
 * console.log(await dynamicPrompt.format({ adjective: "rainy" }));
 *
 * ```
 */
export class FewShotPromptTemplate
  extends BaseStringPromptTemplate
  implements FewShotPromptTemplateInput
{
  lc_serializable = false;

  examples?: InputValues[];

  exampleSelector?: BaseExampleSelector | undefined;

  examplePrompt: PromptTemplate;

  suffix = "";

  exampleSeparator = "\n\n";

  prefix = "";

  templateFormat: TemplateFormat = "f-string";

  validateTemplate = true;

  constructor(input: FewShotPromptTemplateInput) {
    super(input);
    Object.assign(this, input);

    if (this.examples !== undefined && this.exampleSelector !== undefined) {
      throw new Error(
        "Only one of 'examples' and 'example_selector' should be provided"
      );
    }

    if (this.examples === undefined && this.exampleSelector === undefined) {
      throw new Error(
        "One of 'examples' and 'example_selector' should be provided"
      );
    }

    if (this.validateTemplate) {
      let totalInputVariables: string[] = this.inputVariables;
      if (this.partialVariables) {
        totalInputVariables = totalInputVariables.concat(
          Object.keys(this.partialVariables)
        );
      }
      checkValidTemplate(
        this.prefix + this.suffix,
        this.templateFormat,
        totalInputVariables
      );
    }
  }

  _getPromptType(): "few_shot" {
    return "few_shot";
  }

  static lc_name() {
    return "FewShotPromptTemplate";
  }

  private async getExamples(
    inputVariables: InputValues
  ): Promise<InputValues[]> {
    if (this.examples !== undefined) {
      return this.examples;
    }
    if (this.exampleSelector !== undefined) {
      return this.exampleSelector.selectExamples(inputVariables);
    }

    throw new Error(
      "One of 'examples' and 'example_selector' should be provided"
    );
  }

  async partial<NewPartialVariableName extends string>(
    values: PartialValues<NewPartialVariableName>
  ) {
    const newInputVariables = this.inputVariables.filter(
      (iv) => !(iv in values)
    );
    const newPartialVariables = {
      ...(this.partialVariables ?? {}),
      ...values,
    };
    const promptDict = {
      ...this,
      inputVariables: newInputVariables,
      partialVariables: newPartialVariables,
    };
    return new FewShotPromptTemplate(promptDict);
  }

  /**
   * Formats the prompt with the given values.
   * @param values The values to format the prompt with.
   * @returns A promise that resolves to a string representing the formatted prompt.
   */
  async format(values: InputValues): Promise<string> {
    const allValues = await this.mergePartialAndUserVariables(values);
    const examples = await this.getExamples(allValues);

    const exampleStrings = await Promise.all(
      examples.map((example) => this.examplePrompt.format(example))
    );
    const template = [this.prefix, ...exampleStrings, this.suffix].join(
      this.exampleSeparator
    );
    return renderTemplate(template, this.templateFormat, allValues);
  }

  serialize(): SerializedFewShotTemplate {
    if (this.exampleSelector || !this.examples) {
      throw new Error(
        "Serializing an example selector is not currently supported"
      );
    }
    if (this.outputParser !== undefined) {
      throw new Error(
        "Serializing an output parser is not currently supported"
      );
    }
    return {
      _type: this._getPromptType(),
      input_variables: this.inputVariables,
      example_prompt: this.examplePrompt.serialize(),
      example_separator: this.exampleSeparator,
      suffix: this.suffix,
      prefix: this.prefix,
      template_format: this.templateFormat,
      examples: this.examples,
    };
  }

  static async deserialize(
    data: SerializedFewShotTemplate
  ): Promise<FewShotPromptTemplate> {
    const { example_prompt } = data;
    if (!example_prompt) {
      throw new Error("Missing example prompt");
    }
    const examplePrompt = await PromptTemplate.deserialize(example_prompt);

    let examples: Example[];

    if (Array.isArray(data.examples)) {
      examples = data.examples;
    } else {
      throw new Error(
        "Invalid examples format. Only list or string are supported."
      );
    }

    return new FewShotPromptTemplate({
      inputVariables: data.input_variables,
      examplePrompt,
      examples,
      exampleSeparator: data.example_separator,
      prefix: data.prefix,
      suffix: data.suffix,
      templateFormat: data.template_format,
    });
  }
}

export interface FewShotChatMessagePromptTemplateInput
  extends BasePromptTemplateInput<InputValues> {
  /**
   * Examples to format into the prompt. Exactly one of this or
   * {@link exampleSelector} must be
   * provided.
   */
  examples?: Example[];

  /**
   * An {@link BaseMessagePromptTemplate} | {@link BaseChatPromptTemplate} used to format a single example.
   */
  examplePrompt: BaseMessagePromptTemplate | BaseChatPromptTemplate;

  /**
   * String separator used to join the prefix, the examples, and suffix.
   *
   * @defaultValue `"\n\n"`
   */
  exampleSeparator?: string;

  /**
   * An {@link BaseExampleSelector} Examples to format into the prompt. Exactly one of this or
   * {@link examples} must be
   * provided.
   */
  exampleSelector?: BaseExampleSelector | undefined;

  /**
   * A prompt template string to put before the examples.
   *
   * @defaultValue `""`
   */
  prefix?: string;

  /**
   * A prompt template string to put after the examples.
   *
   * @defaultValue `""`
   */
  suffix?: string;

  /**
   * The format of the prompt template. Options are: 'f-string'
   *
   * @defaultValue `f-string`
   */
  templateFormat?: TemplateFormat;

  /**
   * Whether or not to try validating the template on initialization.
   *
   * @defaultValue `true`
   */
  validateTemplate?: boolean;
}