diff --git a/chat/src/agent/agent_handler.py b/chat/src/agent/agent_handler.py index 8b5fe51..5065c40 100644 --- a/chat/src/agent/agent_handler.py +++ b/chat/src/agent/agent_handler.py @@ -52,7 +52,8 @@ def on_tool_end(self, output: ToolMessage, **kwargs: Dict[str, Any]): pass case "search": try: - docs: List[Dict[str, Any]] = [doc.metadata for doc in output.artifact] + result_fields = ("id", "title", "visibility", "work_type", "thumbnail") + docs: List[Dict[str, Any]] = [{k: doc.metadata.get(k) for k in result_fields} for doc in output.artifact] self.socket.send({"type": "search_result", "ref": self.ref, "message": docs}) except json.decoder.JSONDecodeError as e: print(f"Invalid json ({e}) returned from {output.name} tool: {output.content}") diff --git a/chat/src/agent/search_agent.py b/chat/src/agent/search_agent.py index c7211ff..16b41fe 100644 --- a/chat/src/agent/search_agent.py +++ b/chat/src/agent/search_agent.py @@ -8,21 +8,21 @@ from langchain_core.messages.system import SystemMessage from langgraph.graph import END, START, StateGraph, MessagesState from langgraph.prebuilt import ToolNode -from setup import openai_chat_client +from setup import chat_client # Keep your answer concise and keep reading time under 45 seconds. system_message = """ Please provide a brief answer to the question using the tools provided. Include specific details from multiple documents that support your answer. Answer in raw markdown, but not within a code block. When citing source documents, construct Markdown -links using the document's canonical_link field. +links using the document's canonical_link field. Do not include intermediate messages explaining your process. """ tools = [discover_fields, search, aggregate] tool_node = ToolNode(tools) -model = openai_chat_client(streaming=True).bind_tools(tools) +model = chat_client(streaming=True).bind_tools(tools) # Define the function that determines whether to continue or not def should_continue(state: MessagesState) -> Literal["tools", END]: @@ -38,7 +38,7 @@ def should_continue(state: MessagesState) -> Literal["tools", END]: # Define the function that calls the model def call_model(state: MessagesState): messages = [SystemMessage(content=system_message)] + state["messages"] - response: BaseMessage = model.invoke(messages, model=os.getenv("AZURE_OPENAI_LLM_DEPLOYMENT_ID")) + response: BaseMessage = model.invoke(messages) # , model=os.getenv("AZURE_OPENAI_LLM_DEPLOYMENT_ID") # We return a list, because this will get added to the existing list # if socket is not none and the response content is not an empty string return {"messages": [response]} diff --git a/chat/src/handlers/chat.py b/chat/src/handlers/chat.py index 3db7b6b..0b563f1 100644 --- a/chat/src/handlers/chat.py +++ b/chat/src/handlers/chat.py @@ -19,11 +19,9 @@ "debug": [ "answer", "attributes", - "azure_endpoint", "deployment_name", "is_superuser", "k", - "openai_api_version", "prompt", "question", "ref", @@ -36,7 +34,6 @@ "deployment_name", "is_superuser", "k", - "openai_api_version", "prompt", "question", "ref", @@ -79,8 +76,8 @@ def handler(event, context): try: search_agent.invoke( {"messages": [HumanMessage(content=config.question)]}, - config={"configurable": {"thread_id": config.ref}, "callbacks": callbacks, "metadata": {"model_deployment": os.getenv("AZURE_OPENAI_LLM_DEPLOYMENT_ID")}}, - debug=False + config={"configurable": {"thread_id": config.ref}, "callbacks": callbacks}, + debug=True ) except Exception as e: print(f"Error: {e}") diff --git a/chat/src/helpers/metrics.py b/chat/src/helpers/metrics.py index f00abc0..0265fbc 100644 --- a/chat/src/helpers/metrics.py +++ b/chat/src/helpers/metrics.py @@ -6,12 +6,10 @@ def debug_response(config, response, original_question): return { "answer": response, "attributes": config.attributes, - "azure_endpoint": config.azure_endpoint, "deployment_name": config.deployment_name, "is_dev_team": config.api_token.is_dev_team(), "is_superuser": config.api_token.is_superuser(), "k": config.k, - "openai_api_version": config.openai_api_version, "prompt": config.prompt_text, "question": config.question, "ref": config.ref, diff --git a/chat/src/secrets.py b/chat/src/secrets.py index 9e3f0e9..7cf299a 100644 --- a/chat/src/secrets.py +++ b/chat/src/secrets.py @@ -7,16 +7,12 @@ def load_secrets(): EnvironmentMap = [ ['API_TOKEN_SECRET', 'dcapi', 'api_token_secret'], ['OPENSEARCH_ENDPOINT', 'index', 'endpoint'], - ['OPENSEARCH_MODEL_ID', 'index', 'embedding_model'], - ['AZURE_OPENAI_API_KEY', 'azure_openai', 'api_key'], - ['AZURE_OPENAI_LLM_DEPLOYMENT_ID', 'azure_openai', 'llm_deployment_id'], - ['AZURE_OPENAI_RESOURCE_NAME', 'azure_openai', 'resource_name'] + ['OPENSEARCH_MODEL_ID', 'index', 'embedding_model'] ] client = boto3.client("secretsmanager") response = client.batch_get_secret_value(SecretIdList=[ f'{SecretsPath}/infrastructure/index', - f'{SecretsPath}/infrastructure/azure_openai', f'{SecretsPath}/config/dcapi' ]) diff --git a/chat/src/setup.py b/chat/src/setup.py index 3ff3adc..dcc88f7 100644 --- a/chat/src/setup.py +++ b/chat/src/setup.py @@ -1,4 +1,4 @@ -from langchain_openai import AzureChatOpenAI +from langchain_aws import ChatBedrock from handlers.opensearch_neural_search import OpenSearchNeuralSearch from opensearchpy import OpenSearch, RequestsHttpConnection from requests_aws4auth import AWS4Auth @@ -13,11 +13,9 @@ def prefix(value): return "-".join(filter(None, [env_prefix, value])) -def openai_chat_client(**kwargs): - return AzureChatOpenAI( - openai_api_key=os.getenv("AZURE_OPENAI_API_KEY"), - api_version=os.getenv("AZURE_API_VERSION", "2024-08-01-preview"), - azure_endpoint=f"https://{os.getenv("AZURE_OPENAI_RESOURCE_NAME")}.openai.azure.com", +def chat_client(**kwargs): + return ChatBedrock( + model="us.anthropic.claude-3-5-sonnet-20241022-v2:0", **kwargs, ) diff --git a/chat/template.yaml b/chat/template.yaml index 3a26803..876e68c 100644 --- a/chat/template.yaml +++ b/chat/template.yaml @@ -251,6 +251,12 @@ Resources: Resource: - !Sub "arn:aws:s3:::${CheckpointBucket}" - !Sub "arn:aws:s3:::${CheckpointBucket}/*" + - Statement: + - Effect: Allow + Action: + - bedrock:InvokeModel + - bedrock:InvokeModelWithResponseStream + Resource: "*" #* Metadata: #* BuildMethod: nodejs20.x # ChatSyncFunction: