milanchndr/Supply-Chain-AI-Agent
0
1o
2 �$6h;! � @ s� d dl mZmZmZmZ d dlmZmZ d dlm Z m3Z4m Z m Z m
Z
d dlZd dlZd dlmZ ddlmZmZmZ ddlmZ dd lmZ dd5lmZ ee �Z ee d d d
iid� ee jd< e e �Z da da da!da"dd� Z#e#� d*dd�Z$e �%d�dd� �Z&e �%d�dd� �Z'e j%ddgd �d d � �Z(e j%d dgd �e6� d d!� ��Z)e d"kr�e�*d#� ej+�,d$�d%kr� e j-d&d'd(d&d)� dS dS )+� )�Flask�request�jsonify�json)�
create_engine�text)�7JWTManager� jwt_required�get_jwt_identity�create_access_token�get_jwtN)�CORS� )� DATABASE_URL�ROLES_PERMISSIONS�JWT_SECRET_KEY)�setup_qa_chain)�logger)�"create_supply_chain_agent_executorz/*Zoriginszhttp://localhost:8080)Z resourcesr c 8 C sx t �d� zt �d� t� atjjjat �d� W n ty5 } zt j d| � �dd� W Y d } ~ nd } ~ ww z/t �dt9� d�d � �� t t10d11d d �a
t
�� �
}t �d
� W d � n1 s_w Y W n ty� } zt j d| � �dd� W Y d } ~ nd } ~ ww tr�tr�t
r�ztttt
d�at �d� W d S ty� } zt j d| � �dd� W Y d } ~ d S d } ~ ww t � d� d S )Nz!Initializing Flask application...z+Setting up QA chain (for DocumentQATool)...z9QA chain and Bedrock LLM instance for agent tools set up.z(Critical Error setting up QA chain/LLM: T��exc_infoz Setting up database engine: �@������ �12 )Z pool_sizeZ max_overflowz+Database connection to Supabase successful.z+Critical Error setting up database engine: )Z llm_instanceZqa_chain_instanceZdb_engine_instancez(Agent Executor initialized successfully.z(Critical Error creating Agent Executor: z9Cannot create Agent Executor due to missing dependencies.)r �infor �qa_chainZcombine_documents_chainZ llm_chainZllm�bedrock_llm_instance� Exception�errorr �splitr � db_engine�connectr �agent_executor)�e�13connection� r& � D:\Codes\AI Agent\backend\app.py�initialize_app sD 141516 ��17 �� �� �"��r( c C s� z,t �� � }|�td�| |||||d�� |�� W d � W d S 1 s%w Y W d S tyG } zt�d|� �� W Y d }~d S d }~ww )Nz�18 INSERT INTO audit_logs (user_id, role, region, query, success, error_message, timestamp)19 VALUES (:user_id, :role, :region, :query, :success, :error_message, CURRENT_TIMESTAMP)20 )�user_id�role�region�query�success�
error_messagezFailed to log audit: )r! r" �executer Zcommitr r r )r) r* r+ r, r- r. r% r$ r& r& r' � log_audit? s$ 21��22&� ��r0 �/c C s dS )Nz<AI Supply Chain Agent Backend (Multi-Tool Agent) is running.r& r& r& r& r' �homeW s r2 z/healthc C s^ t d utd utd utd ud�} t d u r%d| d< t�d| � �� t| �dfS d| d< t| �dfS )N)Zagent_executor_initializedZ qa_chain_for_tool_initializedZ db_engine_for_tool_initializedZllm_for_agent_initializedZAGENT_NOT_INITIALIZED�statusz-Health check: Agent not initialized. Status: �� ZOK�� )r# r r! r r Zwarningr )r3 r& r& r' � health_check[ s � r6 z/loginZPOST)�methodsc
23 C s| t j} | �d�}| �d�}|r|stddi�dfS z}t�� �n}|�td�d|i��� }|s=tddi�dfW d � W S |\}}}}} t �24|� d �|� d ��s`tddi�dfW d � W S d25|| d �d �}26t t
|�|27d
�} t�d|� d|� d�� td| i�dfW d � W S 1 s�w Y W d S ty� } ztjd|� d| � �dd� tddi�dfW Y d } ~ S d } ~ ww )N�username�passwordr z"Username and password are required� z[SELECT user_id, username, password_hash, role, region FROM users WHERE username = :usernamez Invalid username or passwordi� zutf-8Z
authenticated)r* r+ )r* � app_metadata)Zidentity�additional_claimszUser z (user_id: z+) logged in. JWT created with app_metadata.� access_tokenr5 z Error during login for username �: Tr z"Login failed due to a server error�� )r r �getr r! r" r/ r Zfetchone�bcryptZcheckpw�encoder �strr r r r )
�datar8 r9 r% �resultZuser_id_from_db�_Z
password_hashZ role_from_dbZregion_from_dbr< r= r$ r& r& r' �logini sJ 282930��� ��� �(� ��rG z/queryc C s� t st�d� tddi�dfS t� } t� }|�di ��d�}|�di ��d�}|r-|tvrDd|� d �}t| ||d31d |� td|i�d fS t j32}|�d
d�� � }|sdd}t| |||d |� td|i�dfS t� d| � d|� d|� d|� �� z#t �
|||t33�|�d��}|�dd�}t| |||d� t|dd��W S ty� } zJdt| �� �}d t| �v s�d t| ��� v r�d t| �� �}d }34nd }35t| |||d |� tjd |� d!| � d"t| �� �dd#� td|i�|36fW Y d } ~ S d } ~ ww )$Nz5Agent executor not initialized. Cannot process query.r z8Agent services not fully initialized. Please check logs.r4 r; r* r+ zInvalid application role: z# or role not configured in backend.zN/AFi� �question� zNo question providedr: z Received agent query from user z (app_role: z, app_region: z): )�inputZ user_roleZ user_regionZjwt_claims_for_db�outputz)Agent could not determine a final answer.TZagent_multi_tool)Zanswer�typezError processing query: zUnauthorized accesszpermission deniedzCYou are not authorized to perform this action or access this data. r? z Error processing agent query 'z ' for user r> r )r# r r r r37 r r@ r r0 r r �stripr Zinvoke�dumpsr rC �lower) Zuser_identity_subZ jwt_payloadZapp_roleZ38app_regionZ error_msgrD rH ZresponseZ final_answerr$ Z status_coder& r& r' �handle_agent_query� sN 39 "� $ ��rP �__main__z9Starting Flask development server for Multi-Tool Agent...ZWERKZEUG_RUN_MAIN�trueTz0.0.0.0i� )�debugZhostZportZ use_reloader)N).Zflaskr r r r Z40sqlalchemyr r Zflask_jwt_extendedr r r41 r r �osrA Z42flask_corsr
Zconfigr r r Zutils.langchain_setupr Z
logger_configr Zutils.agent_handlerr �__name__ZappZjwtr r! r r# r( r0 Zrouter2 r6 rG rP r �environr@ �runr& r& r& r' �<module> sH 43 444546
47, 448�