mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-10-01 12:19:39 +02:00
docs: add module docstrings
This commit is contained in:
1 parent
c78f4ced91
commit
d41039773e
1 file changed
+22
-12
@@ -1,4 +1,4 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import os
|
import os
|
||||||
import uuid as U
|
import uuid as U
|
||||||
|
|
||||||
@@ -10,35 +10,47 @@ AUTH_TOKEN: str = os.getenv("AS_TOKEN", "gh0-5f4a8ed2-37c6-4bd7-a0cf-7070eae8115
|
|||||||
|
|
||||||
|
|
||||||
class Module:
|
class Module:
|
||||||
|
""":class:`Module` for fetching prompts from a remote API and submitting LLM outputs.
|
||||||
|
|
||||||
|
Retrieves batches of prompts from a configurable HTTP endpoint (default:
|
||||||
|
``https://mcp.metaheuristic.co/infer``), optionally posts each prompt to an
|
||||||
|
LLM proxy for evaluation, and yields the raw prompts for downstream guard
|
||||||
|
processing.
|
||||||
|
|
||||||
|
Configuration is passed via the ``opts`` dict:
|
||||||
|
|
||||||
|
* ``port`` (int): LLM proxy port. Defaults to ``8718``.
|
||||||
|
* ``max_prompts`` (int): Total prompts to process. Defaults to ``2000``.
|
||||||
|
* ``batch_size`` (int): Prompts fetched per API call. Defaults to ``500``.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
tools_inbox: Async queue that receives guard evaluation results.
|
||||||
|
opts: Module configuration dictionary.
|
||||||
|
prompt_groups: Passed through but not used directly by this module.
|
||||||
|
"""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self, prompt_groups: list[str], tools_inbox: asyncio.Queue, opts: dict = {}
|
self, prompt_groups: list[str], tools_inbox: asyncio.Queue, opts: dict = {}
|
||||||
):
|
):
|
||||||
self.tools_inbox = tools_inbox
|
self.tools_inbox = tools_inbox
|
||||||
self.opts = opts
|
self.opts = opts
|
||||||
self.prompt_groups = prompt_groups
|
self.prompt_groups = prompt_groups
|
||||||
self.max_prompts = self.opts.get("max_prompts", 2000) # Default max M prompts
|
self.max_prompts = self.opts.get("max_prompts", 2000)
|
||||||
self.run_id = U.uuid4().hex
|
self.run_id = U.uuid4().hex
|
||||||
self.batch_size = self.opts.get("batch_size", 500)
|
self.batch_size = self.opts.get("batch_size", 500)
|
||||||
|
|
||||||
async def apply(self):
|
async def apply(self):
|
||||||
for _ in range(max(self.max_prompts // self.batch_size, 1)):
|
for _ in range(max(self.max_prompts // self.batch_size, 1)):
|
||||||
# Fetch prompts from the API
|
|
||||||
prompts = await self.fetch_prompts()
|
prompts = await self.fetch_prompts()
|
||||||
|
|
||||||
if not prompts:
|
if not prompts:
|
||||||
logger.error("No prompts retrieved from the API.")
|
logger.error("No prompts retrieved from the API.")
|
||||||
return
|
return
|
||||||
|
|
||||||
logger.info(f"Retrieved {len(prompts)} prompts.")
|
logger.info(f"Retrieved {len(prompts)} prompts.")
|
||||||
|
|
||||||
for i, prompt in enumerate(
|
for i, prompt in enumerate(
|
||||||
prompts[: self.max_prompts]
|
prompts[: self.max_prompts]
|
||||||
): # Limit to max_prompts
|
):
|
||||||
logger.info(f"Processing prompt {i+1}/{len(prompts)}: {prompt}")
|
logger.info(f"Processing prompt {i+1}/{len(prompts)}: {prompt}")
|
||||||
# response = await self.post_prompt(prompt)
|
|
||||||
# logger.info(f"Response: {response}")
|
|
||||||
yield prompt
|
yield prompt
|
||||||
|
|
||||||
while not self.tools_inbox.empty():
|
while not self.tools_inbox.empty():
|
||||||
ref = await self.tools_inbox.get()
|
ref = await self.tools_inbox.get()
|
||||||
message, _, ready = ref["message"], ref["reply"], ref["ready"]
|
message, _, ready = ref["message"], ref["reply"], ref["ready"]
|
||||||
@@ -55,7 +67,6 @@ class Module:
|
|||||||
"max_tokens": 1050,
|
"max_tokens": 1050,
|
||||||
"temperature": 0.7,
|
"temperature": 0.7,
|
||||||
}
|
}
|
||||||
|
|
||||||
async with httpx.AsyncClient() as client:
|
async with httpx.AsyncClient() as client:
|
||||||
try:
|
try:
|
||||||
response = await client.post(uri, headers=headers, json=data)
|
response = await client.post(uri, headers=headers, json=data)
|
||||||
@@ -71,7 +82,6 @@ class Module:
|
|||||||
"Authorization": f"Bearer {AUTH_TOKEN}",
|
"Authorization": f"Bearer {AUTH_TOKEN}",
|
||||||
"Content-Type": "application/json",
|
"Content-Type": "application/json",
|
||||||
}
|
}
|
||||||
|
|
||||||
async with httpx.AsyncClient() as client:
|
async with httpx.AsyncClient() as client:
|
||||||
try:
|
try:
|
||||||
response = await client.post(
|
response = await client.post(
|
||||||
|
|||||||
Reference in new issue
Block a user