diff --git a/Makefile b/Makefile index 462d72b..ab02702 100644 --- a/Makefile +++ b/Makefile @@ -1,46 +1,52 @@ # --------------------------------------------------------- # Misc. # --------------------------------------------------------- - + # Set the default goal. .DEFAULT_GOAL := build - + # Tell Docker to build images in parallel. COMPOSE_BAKE := true - -# Set the environment variables file to ".env" if an argument is not provided. -ENV_FILE ?= .env -include $(ENV_FILE) - + # Set the Docker Compose profile to "all" if an argument is not provided. DOCKER_COMPOSE_PROFILE ?= all - + # --------------------------------------------------------- # Build the containers. # --------------------------------------------------------- - + .PHONY: build .SILENT: build - -build: - docker compose --profile $(DOCKER_COMPOSE_PROFILE) --env-file $(ENV_FILE) build --no-cache - + +build: + docker compose --profile $(DOCKER_COMPOSE_PROFILE) build + # --------------------------------------------------------- # Start the containers. # --------------------------------------------------------- - + .PHONY: start .SILENT: start - + start: - docker compose --profile $(DOCKER_COMPOSE_PROFILE) --env-file $(ENV_FILE) up -d - + docker compose --profile $(DOCKER_COMPOSE_PROFILE) up -d + # --------------------------------------------------------- # Stop the containers. # --------------------------------------------------------- - + .PHONY: stop .SILENT: stop - -stop: - docker compose --profile $(DOCKER_COMPOSE_PROFILE) --env-file $(ENV_FILE) down + +stop: + docker compose --profile $(DOCKER_COMPOSE_PROFILE) down + +# --------------------------------------------------------- +# Test the agent. +# --------------------------------------------------------- + +.PHONY: test/agent +.SILENT: test/agent + +test/agent: + curl -X POST localhost:8181/api/v1/ -d '{"message":"hello"}'; echo diff --git a/agent/aoai_gpt-4o/Dockerfile b/agent/aoai_gpt-4o/Dockerfile new file mode 100644 index 0000000..d9a60b0 --- /dev/null +++ b/agent/aoai_gpt-4o/Dockerfile @@ -0,0 +1,14 @@ +FROM alpine:3.23 +LABEL image.authors="Vic Fernandez III <@cyberphor>" +WORKDIR /home/swordfish/ +COPY requirements.txt requirements.txt +RUN apk add --no-cache --update python3 py3-pip uvicorn &&\ + pip install --break-system-packages -r requirements.txt &&\ + adduser -D swordfish -h /home/swordfish +COPY entrypoint.sh entrypoint.sh +COPY swordfish/ swordfish/ +RUN chown -R swordfish:swordfish /home/swordfish &&\ + chmod +x entrypoint.sh +USER swordfish +EXPOSE 8181 +ENTRYPOINT [ "./entrypoint.sh" ] diff --git a/agent/aoai_gpt-4o/entrypoint.sh b/agent/aoai_gpt-4o/entrypoint.sh new file mode 100644 index 0000000..ac984f1 --- /dev/null +++ b/agent/aoai_gpt-4o/entrypoint.sh @@ -0,0 +1,4 @@ +#!/bin/sh + +set -e +uvicorn swordfish.main:api --host 0.0.0.0 --port 8181 diff --git a/agent/aoai_gpt-4o/requirements.txt b/agent/aoai_gpt-4o/requirements.txt new file mode 100644 index 0000000..7cddd62 --- /dev/null +++ b/agent/aoai_gpt-4o/requirements.txt @@ -0,0 +1,4 @@ +azure-identity +fastapi +openai +openai-agents diff --git a/agent/aoai_gpt-4o/swordfish/__init__.py b/agent/aoai_gpt-4o/swordfish/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/agent/aoai_gpt-4o/swordfish/main.py b/agent/aoai_gpt-4o/swordfish/main.py new file mode 100644 index 0000000..0b094a3 --- /dev/null +++ b/agent/aoai_gpt-4o/swordfish/main.py @@ -0,0 +1,49 @@ +from os import getenv + +from agents import Agent, OpenAIChatCompletionsModel, Runner, set_tracing_disabled +from agents.mcp import MCPServerStreamableHttp +from azure.identity import ( + AzureAuthorityHosts, + DefaultAzureCredential, + get_bearer_token_provider, +) +from openai import AsyncAzureOpenAI +from fastapi import FastAPI, Request + +AZURE_TOKEN_SCOPES = getenv("AZURE_TOKEN_SCOPES") +AZURE_OPENAI_ENDPOINT = getenv("AZURE_OPENAI_ENDPOINT") +AZURE_OPENAI_API_VERSION = getenv("AZURE_OPENAI_API_VERSION") +MCP_SERVER_ENDPOINT = getenv("MCP_SERVER_ENDPOINT") +AZURE_OPENAI_DEPLOYMENT = getenv("AZURE_OPENAI_DEPLOYMENT") + +token_provider = get_bearer_token_provider( + DefaultAzureCredential(authority=AzureAuthorityHosts.AZURE_GOVERNMENT), + AZURE_TOKEN_SCOPES +) + +client = AsyncAzureOpenAI( + azure_endpoint=AZURE_OPENAI_ENDPOINT, + api_version=AZURE_OPENAI_API_VERSION, + azure_ad_token_provider=token_provider +) +set_tracing_disabled(disabled=True) + +api = FastAPI() +history = [] + +@api.post("/api/v1/") +async def respond(request: Request): + global history + request = await request.json() + async with MCPServerStreamableHttp(name="Swordfish Tools", params={"url": MCP_SERVER_ENDPOINT}, cache_tools_list=True) as mcp_server: + model = OpenAIChatCompletionsModel( + model=AZURE_OPENAI_DEPLOYMENT, + openai_client=client, + ) + agent = Agent( + name="swordfish", + model=model, + mcp_servers=[mcp_server], + ) + result = await Runner.run(starting_agent=agent, input=request["message"]) + return result.final_output diff --git a/agent/openai_gpt-4o/requirements.txt b/agent/openai_gpt-4o/requirements.txt index fdc00a9..0fd0d1e 100644 --- a/agent/openai_gpt-4o/requirements.txt +++ b/agent/openai_gpt-4o/requirements.txt @@ -1,4 +1,7 @@ fastapi openai mcp -strands-agents \ No newline at end of file +psycopg +psycopg[pool] +psycopg[binary] +strands-agents diff --git a/agent/openai_gpt-4o/swordfish/history.py b/agent/openai_gpt-4o/swordfish/history.py new file mode 100644 index 0000000..b759263 --- /dev/null +++ b/agent/openai_gpt-4o/swordfish/history.py @@ -0,0 +1,22 @@ +# agent/swordfish/history.py + +from abc import ABC, abstractmethod + +class HistoryStore(ABC): + @abstractmethod + async def get(self, session_id: str) -> list: + pass + + @abstractmethod + async def save(self, session_id: str, messages: list): + pass + +class InMemoryStore(HistoryStore): + def __init__(self): + self.storage = {} + + async def get(self, session_id: str): + return self.storage.get(session_id, []) + + async def save(self, session_id: str, messages: list): + self.storage[session_id] = messages diff --git a/agent/openai_gpt-4o/swordfish/main.py b/agent/openai_gpt-4o/swordfish/main.py index b567162..0f303ee 100644 --- a/agent/openai_gpt-4o/swordfish/main.py +++ b/agent/openai_gpt-4o/swordfish/main.py @@ -1,36 +1,80 @@ -# Standard-library imports. +import json from os import environ +from contextlib import asynccontextmanager -# Third-party imports. from fastapi import FastAPI, Request from mcp.client.sse import sse_client +from psycopg_pool import AsyncConnectionPool + from strands import Agent from strands.models.openai import OpenAIModel from strands.tools.mcp import MCPClient -api = FastAPI() +from .history import HistoryStore, InMemoryStore + + +class PostgresStore(HistoryStore): + def __init__(self, pool: AsyncConnectionPool): + self.pool = pool + + async def get(self, session_id: str) -> list: + async with self.pool.connection() as conn: + async with conn.cursor() as cur: + await cur.execute("SELECT messages FROM chat_sessions WHERE session_id = %s", (session_id,)) + result = await cur.fetchone() + return result[0] if result else [] + + async def save(self, session_id: str, messages: list): + async with self.pool.connection() as conn: + async with conn.cursor() as cur: + await cur.execute(""" + INSERT INTO chat_sessions (session_id, messages) + VALUES (%s, %s) + ON CONFLICT (session_id) DO UPDATE SET messages = EXCLUDED.messages + """, (session_id, json.dumps(messages))) + +@asynccontextmanager +async def lifespan(app: FastAPI): + store_type = environ.get("STORE_TYPE", "memory").lower() + + match store_type: + case "postgres": + db_url = environ["DATABASE_URL"] + pool = AsyncConnectionPool(conninfo=db_url, open=False) + await pool.open() + app.state.history = PostgresStore(pool) + yield + await pool.close() + case _: + app.state.history = InMemoryStore() + yield + +api = FastAPI(lifespan=lifespan) sse_mcp_client = MCPClient(lambda: sse_client("http://swordfish-tools:8282/sse")) -history = [] @api.post("/api/v1/") async def respond(request: Request): - global history - request = await request.json() + store: HistoryStore = request.app.state.history + + body = await request.json() + session_id = body.get("session_id", "default_session") + user_message = body.get("message") + + current_history = await store.get(session_id) + with sse_mcp_client: model = OpenAIModel( - client_args={ - "api_key": environ["OPENAI_API_KEY"], - }, - model_id=request.get("model", "gpt-4o"), - params={ - "max_tokens": 1000, - "temperature": 0.7, - }, + client_args={"api_key": environ["OPENAI_API_KEY"]}, + model_id=body.get("model", "gpt-4o"), + params={"max_tokens": 1000, "temperature": 0.7}, ) tools = sse_mcp_client.list_tools_sync() + agent = Agent(model=model, tools=tools) - agent.messages = history - response = agent(prompt=request["message"]) - history = agent.messages + agent.messages = current_history + response = agent(prompt=user_message) + + await store.save(session_id, agent.messages) + return response.message["content"][0]["text"] diff --git a/compose.yml b/compose.yml index 2892b22..0432332 100644 --- a/compose.yml +++ b/compose.yml @@ -1,14 +1,12 @@ services: - + docs: - profiles: [ all, docs ] + profiles: [ docs ] build: docs image: swordfish/docs:latest container_name: swordfish-docs ports: - "5050:5050" - env_file: - - .env emass: profiles: [ all, emass ] @@ -17,9 +15,7 @@ services: container_name: swordfish-emass ports: - "4010:4010" - env_file: - - .env - + tools: profiles: [ all, tools ] build: tools @@ -28,18 +24,18 @@ services: ports: - "8282:8282" env_file: - - .env - + - tools/.env + agent: profiles: [ all, agent ] - build: agent/openai_gpt-4o + build: agent/aoai_gpt-4o image: swordfish/agent:latest container_name: swordfish-agent ports: - "8181:8181" env_file: - - .env - + - agent/aoai_gpt-4o/.env + frontend: profiles: [ all, frontend ] build: frontend @@ -47,5 +43,3 @@ services: container_name: swordfish-frontend ports: - "80:8501" - env_file: - - .env diff --git a/docs/swordfish/developer-guide/agents.md b/docs/swordfish/developer-guide/agents.md new file mode 100644 index 0000000..9763190 --- /dev/null +++ b/docs/swordfish/developer-guide/agents.md @@ -0,0 +1,16 @@ +# Agents + +## Azure GPT-4o +```bash +export AZURE_TENANT_ID="" +export AZURE_CLIENT_ID="" +export AZURE_CLIENT_SECRET="" + +export AZURE_CLOUD="" +export AZURE_TOKEN_SCOPES="https://cognitiveservices.azure.us/.default" + +export AZURE_OPENAI_ENDPOINT="https://.openai.azure.us/" +export AZURE_OPENAI_API_VERSION="" +export AZURE_OPENAI_DEPLOYMENT="swordfish" +export MCP_SERVER_ENDPOINT="http://swordfish-tools:8282/sse" +``` diff --git a/docs/swordfish/developer-guide/infrastructure.md b/docs/swordfish/developer-guide/infrastructure.md new file mode 100644 index 0000000..4a536c9 --- /dev/null +++ b/docs/swordfish/developer-guide/infrastructure.md @@ -0,0 +1,47 @@ +# Infrastructure + +## Azure + +**Step 1.** Deploy cloud resources using Infrastructure-as-Code. Said resources must ensure there is an Azure OpenAI (AOAI) service with a deployment specific for your project. + +**Step 2.** Activate your Owner role on the AOAI service and then, login to Azure via CLI. +```bash +az login --use-device-code +``` + +**Step 3.** Get the client ID and client secret for your Service Principal and then add them as environment variables to your project. + +**Step 4.** Get the object ID of service principal. +```bash +export OBJECT_ID=$(az ad sp show --id $CLIENT_ID --query id -o tsv) +``` + +**Step 5.** Get the resource ID of Azure OpenAI service. +```bash +export RESOURCE_ID=$(az cognitiveservices account show \ + --name $AOAI_SERVICE_NAME \ + --resource-group $AOAI_SERVICE_RESOURCE_GROUP \ + --query id -o tsv) +``` + +**Step 6.** Assign the service principal the "Cognitive Services OpenAI User" role on the object ID. +```bash +az role assignment create \ + --assignee $OBJECT_ID \ + --role "Cognitive Services OpenAI User" \ + --scope $RESOURCE_ID +``` + +**Step 7.** Assign the service principal the "Cognitive Services OpenAI Contributor" role on the object ID. +```bash +az role assignment create \ + --assignee $OBJECT_ID \ + --role "Cognitive Services OpenAI Contributor" \ + --scope $RESOURCE_ID +``` + +**Step 8.** Run your project. Running your project should include executing code that authenticates with the AOAI service as your Service Principal. + +## References +* [Compare Azure Government and global Azure](https://learn.microsoft.com/en-us/azure/azure-government/compare-azure-government-global-azure) +* [How to switch between OpenAI and Azure OpenAI endpoints](https://learn.microsoft.com/en-us/azure/developer/ai/how-to/switching-endpoints?tabs=azure-openai&pivots=python) diff --git a/frontend/swordfish/main.py b/frontend/swordfish/main.py index b234505..864bae2 100644 --- a/frontend/swordfish/main.py +++ b/frontend/swordfish/main.py @@ -1,8 +1,12 @@ import streamlit as st import requests +import uuid st.title("Swordfish") +if "session_id" not in st.session_state: + st.session_state.session_id = str(uuid.uuid4()) + if "messages" not in st.session_state: st.session_state.messages = [] @@ -10,7 +14,24 @@ with st.chat_message(message["role"]): st.markdown(message["content"]) -if prompt := st.chat_input("Say something"): +PROMPT_SUGGESTIONS = [ + "draft a poam entry for the latest finding", + "summarize the risk posture of my oldest information system", + "is this cyber tasking order applicable to any of my emass records?", +] + +selected_prompt = None + +if not st.session_state.messages: + cols = st.columns(len(PROMPT_SUGGESTIONS)) + for col, chip in zip(cols, PROMPT_SUGGESTIONS): + if col.button(chip, use_container_width=True): + selected_prompt = chip + +chat_input = st.chat_input("Ask Swordfish...") +prompt = chat_input or selected_prompt + +if prompt: st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user"): @@ -18,9 +39,13 @@ response = requests.post( "http://swordfish-agent:8181/api/v1/", - json={"message": prompt}, + json={ + "message": prompt, + "session_id": st.session_state.session_id + }, timeout=60, ) + response.raise_for_status() text = response.json() @@ -28,3 +53,6 @@ with st.chat_message("assistant"): st.markdown(text) + + if selected_prompt: + st.rerun() diff --git a/tools/requirements.txt b/tools/requirements.txt index 04d899a..285c46f 100644 --- a/tools/requirements.txt +++ b/tools/requirements.txt @@ -1,2 +1 @@ -mcp[cli] -strands-agents-tools +fastmcp