Using Azure OpenAI Bring Your Data with Langchain

[hegeduscs]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

extra_body={
    "data_sources":[
        {
            "type": "azure_search",
            "parameters": {
                    "endpoint": AZURE_AISEARCH_ENDPOINT,
                    "index_name": AZURE_AISEARCH_INDEX_NAME,
                    "authentication": {
                        "type": "api_key",
                        "key": AZURE_AISEARCH_API_KEY,
                    },
                    "query_type": "vector_semantic_hybrid", # OPTIONS: simple, semantic, vector, vector_simple_hybrid, vector_semantic_hybrid
                    "fields_mapping": {
                        "content_fields" : ["content"],
                        "title_field": "Header1", # document title 
                        "vector_fields": ["content_vector"],
                        "filepath_field": "",
                        "url_field" : "page" #page number

                    },
                    "embedding_dependency" :{
                        "deployment_name": AZURE_OPENAI_EMBEDDING_MODEL,
                        "type": "deployment_name", 
                    },
                    "semantic_configuration": "semantic-config",
            }
        }   
    ]
}

from langchain_community.chat_message_histories import SQLChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, HumanMessagePromptTemplate
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import AzureChatOpenAI

llm = AzureChatOpenAI(
    azure_endpoint=AZURE_OPENAI_ENDPOINT,
    api_key=AZURE_OPENAI_API_KEY,
    api_version="2024-10-21",
    deployment_name="gpt-4o-mini",
    extra_body=extra_body
)

SYSTEM_PROMPT_RAW = """
You are a helpful assistant.
    """

prompt_with_history = ChatPromptTemplate.from_messages(
    [
        SystemMessage(SYSTEM_PROMPT_RAW),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}")
    ]
)

chain_without_history = prompt_with_history | llm

def get_session_history(session_id):
    return SQLChatMessageHistory(session_id, "sqlite:///memory.db")

chain_with_history = RunnableWithMessageHistory(
    chain_without_history,
    get_session_history,
    history_messages_key = "history",
    input_messages_key="input",
)

response_1 = chain_with_history.invoke(
    {"input":"What is something?"},
    config={"configurable": {"session_id": "1"}},
)
print(response_1)

Description

I am trying to use Azure OpenAI's Bring Your Data feature ( https://learn.microsoft.com/en-us/azure/ai-services/openai/use-your-data-quickstart?tabs=command-line%2Cjavascript-keyless%2Ctypescript-keyless%2Cpython-new&pivots=programming-language-python ) in a Langchain chain.

I am properly able to configure Azure AI Search as the source,but I am not finding any way to get back the citation part of the response through Langchain. With simple OpenAI client it works well, and this is what returned in the background:

{
  "id": "",
  "choices": [
    {
      "finish_reason": "stop",
      "index": 0,
      "logprobs": null,
      "message": {
        "content": "Response [doc1].",
        "refusal": null,
        "role": "assistant",
        "audio": null,
        "function_call": null,
        "tool_calls": null,
        "end_turn": true,
        **"context": {
          "citations": [
            {
              "content": "chunk",
              "title": "",
              "url": "13",
              "filepath": "",
              "chunk_id": "0"
            },**
            {
...
    "completion_tokens_details": null,
    "prompt_tokens_details": null
  }
}

System Info

System Information

OS: Windows
OS Version: 10.0.22631
Python Version: 3.12.0 (tags/v3.12.0:0fb18b0, Oct 2 2023, 13:03:39) [MSC v.1935 64 bit (AMD64)]

Package Information

langchain_core: 0.3.15
langchain: 0.3.6
langchain_community: 0.3.3
langsmith: 0.1.136
langchain_experimental: 0.3.2
langchain_openai: 0.2.5
langchain_text_splitters: 0.3.1

Optional packages not installed

langgraph
langserve

Other Dependencies

aiohttp: 3.10.10
async-timeout: Installed. No version info available.
dataclasses-json: 0.6.7
httpx: 0.27.2
jsonpatch: 1.33
numpy: 1.26.4
openai: 1.52.0
orjson: 3.10.7
packaging: 24.1
pydantic: 2.9.2
pydantic-settings: 2.6.0
PyYAML: 6.0.2
requests: 2.32.3
requests-toolbelt: 1.0.0
SQLAlchemy: 2.0.36
tenacity: 9.0.0
tiktoken: 0.8.0
typing-extensions: 4.12.2

[Krish054]

@hegeduscs Any updates? Were you able to figure out a solution?

[hegeduscs]

Yeah, I forked the langchain-core and edited the openai_base.py and azure.py and manually injecting the missing parts.

I added these 2 lines in azure.py on line 740:

if "context" in response["choices"][0]["message"].keys():
    chat_result.llm_output["model_extra"] = response["choices"][0]["message"]["context"]

and these to line 257 in openai_base:

if _dict.get("context"):
        additional_kwargs["context"] = dict(_dict["context"])

[Krish054]

I use two agents - one for retrieving data from azure search index and the other from a sql table. So, my llm does not need an extra_body param always. But I appreciate the response, thanks a lot!

[zahhar]

@hegeduscs - I was facing the same issue and found your reply while working on the solution. Did you consider proposing a PR to langchain-core to make this feature available for everyone? I would appreciate that a lot!