1# pylint: disable=line-too-long,useless-suppression
2# ------------------------------------
3# Copyright (c) Microsoft Corporation.
4# Licensed under the MIT License.
5# ------------------------------------
6"""
7DESCRIPTION:
8    This sample demonstrates how to add files to agent during the vector store creation.
9 
10USAGE:
11    python sample_agents_enterprise_file_search.py
12 
13    Before running the sample:
14 
15    pip install azure-ai-projects azure-ai-agents azure-identity azure-ai-ml
16 
17    Set these environment variables with your own values:
18    1) PROJECT_ENDPOINT - The Azure AI Project endpoint, as found in the Overview
19                          page of your Azure AI Foundry portal.
20    2) MODEL_DEPLOYMENT_NAME - The deployment name of the AI model, as found under the "Name" column in
21       the "Models + endpoints" tab in your Azure AI Foundry project.
22    3) AZURE_BLOB_URI - The URI of the blob storage where the file is uploaded. In the format:
23         azureml://subscriptions/{subscription-id}/resourcegroups/{resource-group-name}/workspaces/{workspace-name}/datastores/{datastore-name}/paths/{path-to-file}
24"""
25 
26import os
27from azure.ai.projects import AIProjectClient
28from azure.ai.agents.models import (
29    FileSearchTool,
30    ListSortOrder,
31    VectorStoreDataSource,
32    VectorStoreDataSourceAssetType,
33    RunStatus,
34)
35from azure.identity import DefaultAzureCredential
36 
37project_client = AIProjectClient(
38    endpoint=os.environ["PROJECT_ENDPOINT"],
39    credential=DefaultAzureCredential(),
40)
41 
42with project_client:
43    agents_client = project_client.agents
44 
45    # [START upload_file_and_create_agent_with_file_search]
46    # If provided, we will upload the local file to Azure and will use it for vector store creation.
47    # Otherwise, we'll use a previously created dataset reference
48    if "AZURE_BLOB_URI" in os.environ:
49        asset_uri = os.environ["AZURE_BLOB_URI"]
50    else:
51        dataset_name = os.environ["AZURE_DATASET_NAME"]
52        dataset_version = os.environ["AZURE_DATASET_VERSION"]
53        dataset = project_client.datasets.get(name=dataset_name, version=dataset_version)
54        asset_uri = dataset.id
55 
56    # Create a vector store and wait for it to be processed
57    ds = VectorStoreDataSource(asset_identifier=asset_uri, asset_type=VectorStoreDataSourceAssetType.URI_ASSET)
58    vector_store = agents_client.vector_stores.create_and_poll(data_sources=[ds], name="sample_vector_store")
59    print(f"Created vector store, vector store ID: {vector_store.id}")
60    vector_store_files = {}
61    for fle in agents_client.vector_store_files.list(vector_store.id):
62        uploaded_file = agents_client.files.get(fle.id)
63        vector_store_files[fle.id] = uploaded_file.filename
64 
65    # Create a file search tool
66    file_search_tool = FileSearchTool(vector_store_ids=[vector_store.id])
67 
68    # Notices that FileSearchTool as tool and tool_resources must be added or the agent unable to search the file
69    agent = agents_client.create_agent(
70        model=os.environ["MODEL_DEPLOYMENT_NAME"],
71        name="my-agent",
72        instructions="You are helpful agent",
73        tools=file_search_tool.definitions,
74        tool_resources=file_search_tool.resources,
75    )
76    # [END upload_file_and_create_agent_with_file_search]
77    print(f"Created agent, agent ID: {agent.id}")
78 
79    thread = agents_client.threads.create()
80    print(f"Created thread, thread ID: {thread.id}")
81 
82    message = agents_client.messages.create(
83        thread_id=thread.id, role="user", content="What is the content of the files you have access to?"
84    )
85    print(f"Created message, message ID: {message.id}")
86 
87    run = agents_client.runs.create_and_process(thread_id=thread.id, agent_id=agent.id)
88    print(f"Created run, run ID: {run.id}")
89 
90    if run.status == RunStatus.FAILED:
91        print(f"Run failed with: |{run.last_error}|")
92 
93    agents_client.vector_stores.delete(vector_store.id)
94    print("Deleted vector store")
95 
96    agents_client.delete_agent(agent.id)
97    print("Deleted agent")
98 
99    messages = agents_client.messages.list(thread_id=thread.id, order=ListSortOrder.ASCENDING)
100    for msg in messages:
101        if msg.text_messages:
102            last_text = msg.text_messages[-1].text.value
103            for annotation in msg.text_messages[-1].text.annotations:
104 
105                citation = vector_store_files.get(annotation.file_citation.file_id, annotation.file_citation.file_id)
106                last_text = last_text.replace(annotation.text, f" [{citation}]")
107            print(f"{msg.role}: {last_text}")