|
2 | 2 | # Licensed under the MIT license. See LICENSE.md file in the project root for full license information. |
3 | 3 |
|
4 | 4 | from typing import Any |
5 | | -import azure.identity.aio |
6 | 5 | from quart import Blueprint, jsonify, request, Response, render_template, current_app |
7 | 6 |
|
8 | 7 | import asyncio |
|
22 | 21 | AgentStreamEvent |
23 | 22 | ) |
24 | 23 |
|
25 | | -from src.quartapp.config_helper import ConfigHelper |
26 | | - |
27 | | -config = ConfigHelper() |
28 | | - |
29 | 24 | bp = Blueprint("chat", __name__, template_folder="templates", static_folder="static") |
30 | 25 |
|
31 | 26 |
|
32 | 27 | @bp.before_app_serving |
33 | | -async def configure_assistant_client(): |
| 28 | +async def start_server(): |
34 | 29 |
|
35 | | - |
36 | 30 | ai_client = AIProjectClient.from_connection_string( |
37 | | - credential=DefaultAzureCredential( |
38 | | - exclude_shared_token_cache_credential=True), |
| 31 | + credential=DefaultAzureCredential(exclude_shared_token_cache_credential=True), |
39 | 32 | conn_str=os.environ["PROJECT_CONNECTION_STRING"], |
40 | 33 | ) |
41 | 34 |
|
42 | | - agent_id = config.get("Agent", "AGENT_ID") |
| 35 | + # TODO: add more files are not supported for citation at the moment |
| 36 | + files = ["product_info_1.md"] |
| 37 | + file_ids = [] |
| 38 | + for file in files: |
| 39 | + file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'files', file)) |
| 40 | + print(f"Uploading file {file_path}") |
| 41 | + file = await ai_client.agents.upload_file_and_poll(file_path=file_path, purpose=FilePurpose.AGENTS) |
| 42 | + file_ids.append(file.id) |
43 | 43 |
|
44 | | - agent = None |
45 | | - |
46 | | - if agent_id: |
47 | | - try: |
48 | | - agent = await ai_client.agents.get_agent(agent_id) |
49 | | - print(f"Agent already exists, agent ID: {agent.id}") |
50 | | - except Exception as e: |
51 | | - print(f"Agent not found: {e}") |
52 | | - |
53 | | - if agent is None: |
54 | | - files = ["product_info_1.md", "product_info_2.md"] |
55 | | - file_ids = [] |
56 | | - for file in files: |
57 | | - file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'files', file)) |
58 | | - print(f"Uploading file {file_path}") |
59 | | - file = await ai_client.agents.upload_file_and_poll(file_path=file_path, purpose=FilePurpose.AGENTS) |
60 | | - file_ids.append(file.id) |
61 | | - |
62 | | - vector_store = await ai_client.agents.create_vector_store(file_ids=file_ids, name="sample_store") |
| 44 | + vector_store = await ai_client.agents.create_vector_store(file_ids=file_ids, name="sample_store") |
63 | 45 |
|
64 | | - file_search_tool = FileSearchTool(vector_store_ids=[vector_store.id]) |
65 | | - |
66 | | - tool_set = AsyncToolSet() |
67 | | - tool_set.add(file_search_tool) |
68 | | - |
69 | | - agent = await ai_client.agents.create_agent( |
70 | | - model="gpt-4-1106-preview", name="my-assistant", instructions="You are helpful assistant", tools = tool_set.definitions, tool_resources=tool_set.resources |
71 | | - ) |
| 46 | + file_search_tool = FileSearchTool(vector_store_ids=[vector_store.id]) |
| 47 | + |
| 48 | + tool_set = AsyncToolSet() |
| 49 | + tool_set.add(file_search_tool) |
| 50 | + |
| 51 | + print(f"ToolResource: {tool_set.resources}") |
72 | 52 |
|
73 | | - print(f"Created agent, agent ID: {agent.id}") |
74 | | - |
75 | | - config.set('Agent', 'AGENT_ID', agent.id) |
76 | | - config.save() |
77 | | - |
| 53 | + agent = await ai_client.agents.create_agent( |
| 54 | + model="gpt-4o-mini", name="my-assistant", instructions="You are helpful assistant", tools = tool_set.definitions, tool_resources=tool_set.resources |
| 55 | + ) |
78 | 56 |
|
| 57 | + print(f"Created agent, agent ID: {agent.id}") |
| 58 | + |
79 | 59 | bp.ai_client = ai_client |
80 | 60 | bp.agent = agent |
| 61 | + bp.vector_store = vector_store |
| 62 | + bp.file_ids = file_ids |
81 | 63 |
|
82 | 64 |
|
83 | 65 | @bp.after_app_serving |
84 | | -async def shutdown_assistant_client(): |
| 66 | +async def stop_server(): |
| 67 | + for file_id in bp.file_ids: |
| 68 | + await bp.ai_client.agents.delete_file(file_id) |
| 69 | + print(f"Deleted file {file_id}") |
| 70 | + |
| 71 | + await bp.ai_client.agents.delete_vector_store(bp.vector_store.id) |
| 72 | + print(f"Deleted vector store {bp.vector_store.id}") |
| 73 | + |
| 74 | + await bp.ai_client.agents.delete_agent(bp.agent.id) |
| 75 | + |
| 76 | + print(f"Deleted agent {bp.agent.id}") |
| 77 | + |
85 | 78 | await bp.ai_client.close() |
| 79 | + print("Closed AIProjectClient") |
86 | 80 |
|
87 | 81 | @bp.get("/") |
88 | 82 | async def index(): |
@@ -160,10 +154,7 @@ async def chat(): |
160 | 154 |
|
161 | 155 | @bp.route('/fetch-document', methods=['GET']) |
162 | 156 | async def fetch_document(): |
163 | | - filename = request.args.get('filename') |
164 | | - current_app.logger.info(f"Fetching document: {filename}") |
165 | | - if not filename: |
166 | | - return jsonify({"error": "Filename is required"}), 400 |
| 157 | + filename = "product_info_1.md" |
167 | 158 |
|
168 | 159 | # Get the file path from the mapping |
169 | 160 | file_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'files', filename)) |
|
0 commit comments