added authentication token for oeg (fedd2b54) · Commits · OOP / code / ai2

ai_agent/routes/groq.py

+9 −4

Changes for ai_agent/routes/groq.py: 9 added lines, 4 removed lines.
Original line number	Diff line number	Diff line
		import logging
		from typing import Annotated

		from fastapi import APIRouter, Body, Query
		from fastapi import APIRouter, Body, Depends, Query
		from fastapi.responses import JSONResponse, StreamingResponse

		from ai_agent.memory import RedisShortTermMemory
		from ai_agent.sub_agents.groq_agent import create_groq_agent
		from ai_agent.utils import stream_agent_response
		from ai_agent.security import get_authorization

		logger = logging.getLogger(__name__)

		@@ -15,14 +16,17 @@ memory_store = RedisShortTermMemory()


		@router.post("/groq-mcp")
		async def groq_query(payload: Annotated[dict \| None, Body()] = None):
		async def groq_query(
		authorization: Annotated[str, Depends(get_authorization)],
		payload: Annotated[dict \| None, Body()] = None,
		):
		query = payload.get("query") if payload else None
		session_id = payload.get("session_id") if payload else None

		if not query:
		return JSONResponse(status_code=400, content={"error": "No query provided"})

		agent = await create_groq_agent()
		agent = await create_groq_agent(authorization)
		try:
		effective_query = await memory_store.augment_query(query, session_id)
		result = await agent.run(effective_query)
		@@ -40,13 +44,14 @@ async def groq_query(payload: Annotated[dict \| None, Body()] = None):

		@router.get("/groq-mcp/stream")
		async def groq_stream(
		authorization: Annotated[str, Depends(get_authorization)],
		query: str \| None = Query(default=None),
		session_id: str \| None = Query(default=None),
		):
		if not query:
		return JSONResponse(status_code=400, content={"error": "No query provided"})

		agent = await create_groq_agent()
		agent = await create_groq_agent(authorization)

		try:
		effective_query = await memory_store.augment_query(query, session_id)

ai_agent/security.py

0 → 100644

+18 −0

Changes for ai_agent/security.py: 18 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		from typing import Annotated

		from fastapi import Header, HTTPException


		def get_authorization(
		authorization: Annotated[str \| None, Header()] = None,
		) -> str:
		if authorization is not None:
		scheme, _, token = authorization.partition(" ")
		if scheme.lower() == "bearer" and token.strip():
		return f"Bearer {token.strip()}"

		raise HTTPException(
		status_code=401,
		detail="Bearer token required",
		headers={"WWW-Authenticate": "Bearer"},
		)
		No newline at end of file

ai_agent/sub_agents/groq_agent.py

+9 −2

Changes for ai_agent/sub_agents/groq_agent.py: 9 added lines, 2 removed lines.
Original line number	Diff line number	Diff line
		@@ -20,7 +20,7 @@ def require_env(name: str) -> str:
		return value.strip()


		async def create_groq_agent():
		async def create_groq_agent(authorization: str):
		"""Create an MCPAgent using Groq backend."""
		try:
		mcp_server_url = require_env("MCP_SERVER_URL")
		@@ -36,7 +36,14 @@ async def create_groq_agent():

		client = MCPClient(
		{
		"mcpServers": {"http": {"url": mcp_server_url}},
		"mcpServers": {
		"http": {
		"url": mcp_server_url,
		"headers": {
		"Authorization": authorization,
		},
		}
		},
		"use-oauth2": False,
		"use-oidc": False,
		}

mcp_module/security/authorization.py

0 → 100644

+13 −0

Changes for mcp_module/security/authorization.py: 13 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		from fastmcp.exceptions import ToolError
		from fastmcp.server.dependencies import get_http_headers


		def get_oeg_authorization_headers() -> dict[str, str]:
		authorization = get_http_headers().get("authorization")

		if authorization is not None:
		scheme, _, token = authorization.partition(" ")
		if scheme.lower() == "bearer" and token.strip():
		return {"Authorization": f"Bearer {token.strip()}"}

		raise ToolError("Bearer token required to call OEG")
		No newline at end of file

mcp_module/tools/edge_application.py

+2 −0

Changes for mcp_module/tools/edge_application.py: 2 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -18,6 +18,7 @@ from mcp_module.models.edge_application_models import (
		InstantiateAppInput,
		RegisterAppDefinitionInput,
		)
		from mcp_module.security.authorization import get_oeg_authorization_headers

		# Configure logging
		logging.basicConfig(
		@@ -35,6 +36,7 @@ load_dotenv(dotenv_path=env_path)

		def get_edge_headers(x_correlator: str \| None = None) -> dict[str, str]:
		headers = {"accept": "application/json"}
		headers.update(get_oeg_authorization_headers())

		if x_correlator is not None:
		if not re.fullmatch(r"[a-zA-Z0-9-_:;.\\/<>{}]{0,256}", x_correlator):