Merge pull request #7 from macrocosm-os/feature/openapi

README.md

@@ -77,7 +77,7 @@ At present, the API provides two endpoints: `/chat` (live) and `/echo` (test).
 
 `/chat` is used to chat with the network and receive a response. It requires a JSON payload structured as per the QueryValidatorParams class.
 The request payload requires the following parameters encapsulated within the [`QueryValidatorParams`](./validators/base.py) data class:
-- `k_miners: int`: The number of miners from which to request responses.
+- `k: int`: The number of miners from which to request responses.
 - `exclude: List[str]`: A list of roles or agents to exclude from querying.
 - `roles: List[str]`: The roles of the agents to query.
 - `messages: List[str]`: The messages to be sent to the network.
@@ -124,6 +124,10 @@ Final JSON:
 ```
 After verifying that the server is responding to requests locally, you can test the server on a remote machine.
 
+## OpenAPI
+To access OpenAPI specification, go to:
+[http://localhost:10000/docs](http://localhost:10000/docs)
+
 ### Troubleshooting
 
 If you do not receive a response from the server, check that the server is running and that the port is open on the server. You can open the port using the following commands:

middlewares.py → common/middlewares.py

@@ -8,6 +8,10 @@ EXPECTED_ACCESS_KEY = os.environ.get("EXPECTED_ACCESS_KEY")
 
 @middleware
 async def api_key_middleware(request: Request, handler):
+    if request.path.startswith("/docs") or request.path.startswith("/static/swagger"):
+        # Skip checks when accessing OpenAPI documentation.
+        return await handler(request)
+
     # Logging the request
     bt.logging.info(f"Handling {request.method} request to {request.path}")
 
@@ -23,6 +27,10 @@ async def api_key_middleware(request: Request, handler):
 
 @middleware
 async def json_parsing_middleware(request: Request, handler):
+    if request.path.startswith("/docs") or request.path.startswith("/static/swagger"):
+        # Skip checks when accessing OpenAPI documentation.
+        return await handler(request)
+
     try:
         # Parsing JSON data from the request
         request["data"] = await request.json()

common/schemas.py

@@ -0,0 +1,29 @@
+from marshmallow import Schema, fields
+
+
+class QueryChatSchema(Schema):
+    k = fields.Int(description="The number of miners from which to request responses.")
+    exclude = fields.List(fields.Str(), description="A list of roles or agents to exclude from querying.")
+    roles = fields.List(fields.Str(), required=True, description="The roles of the agents to query.")
+    messages = fields.List(fields.Str(), required=True, description="The messages to be sent to the network.")
+    timeout = fields.Int(description="The time in seconds to wait for a response.")
+    prefer = fields.Str(description="The preferred response format, can be either 'longest' or 'shortest'.")
+    sampling_mode = fields.Str(
+        description="The mode of sampling to use, defaults to 'random'. Can be either 'random' or 'top_incentive'.")
+
+
+class StreamChunkSchema(Schema):
+    delta = fields.Str(required=True, description="The new chunk of response received.")
+    finish_reason = fields.Str(description="The reason for the response completion, if applicable.")
+    accumulated_chunks = fields.List(fields.Str(), description="All accumulated chunks of responses.")
+    accumulated_chunks_timings = fields.List(fields.Float(), description="Timing for each chunk received.")
+    timestamp = fields.Str(required=True, description="The timestamp at which the chunk was processed.")
+    sequence_number = fields.Int(required=True, description="A sequential identifier for the response part.")
+    selected_uid = fields.Int(required=True, description="The identifier for the selected response source.")
+
+
+class StreamErrorSchema(Schema):
+    error = fields.Str(required=True, description="Description of the error occurred.")
+    timestamp = fields.Str(required=True, description="The timestamp of the error.")
+    sequence_number = fields.Int(required=True, description="A sequential identifier for the error.")
+    finish_reason = fields.Str(default="error", description="Indicates an error completion.")

requirements.txt

@@ -1,3 +1,4 @@
 git+https://github.com/opentensor/prompting.git@features/move-validator-into-prompting
 aiohttp
-deprecated
+deprecated
+aiohttp_apispec>=2.2.3

server.py

@@ -1,15 +1,25 @@
 import asyncio
-import utils
+
 import bittensor as bt
 from aiohttp import web
-from validators import S1ValidatorAPI, QueryValidatorParams, ValidatorAPI
-from middlewares import api_key_middleware, json_parsing_middleware
-
-
+from aiohttp_apispec import docs, request_schema, response_schema, setup_aiohttp_apispec, validation_middleware
+
+from common import utils
+from common.middlewares import api_key_middleware, json_parsing_middleware
+from common.schemas import QueryChatSchema, StreamChunkSchema, StreamErrorSchema
+from validators import QueryValidatorParams, S1ValidatorAPI, ValidatorAPI
+
+
+@docs(
+    tags=["Prompting API"],
+    summary="Chat",
+    description="Chat endpoint."
+)
+@request_schema(QueryChatSchema)
+@response_schema(StreamChunkSchema, 200)
+@response_schema(StreamErrorSchema, 400)
 async def chat(request: web.Request) -> web.StreamResponse:
-    """
-    Chat endpoint for the validator.
-    """
+    """Chat endpoint for the validator"""
     params = QueryValidatorParams.from_request(request)
 
     # Access the validator from the application context
@@ -19,6 +29,14 @@ async def chat(request: web.Request) -> web.StreamResponse:
     return response
 
 
+@docs(
+    tags=["Prompting API"],
+    summary="Echo test",
+    description="Echo endpoint for testing purposes."
+)
+@request_schema(QueryChatSchema)
+@response_schema(StreamChunkSchema, 200)
+@response_schema(StreamErrorSchema, 400)
 async def echo_stream(request: web.Request) -> web.StreamResponse:
     return await utils.echo_stream(request)
 
@@ -27,19 +45,32 @@ class ValidatorApplication(web.Application):
     def __init__(self, validator_instance=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        self["validator"] = (
-            validator_instance if validator_instance else S1ValidatorAPI()
-        )
+        self["validator"] = validator_instance if validator_instance else S1ValidatorAPI()
 
         # Add middlewares to application
-        self.add_routes([web.post("/chat/", chat), web.post("/echo/", echo_stream)])
+        self.add_routes(
+            [
+                web.post("/chat/", chat),
+                web.post("/echo/", echo_stream),
+            ]
+        )
+        self.setup_openapi()
         self.setup_middlewares()
         # TODO: Enable rewarding and other features
 
     def setup_middlewares(self):
+        self.middlewares.append(validation_middleware)
         self.middlewares.append(json_parsing_middleware)
         self.middlewares.append(api_key_middleware)
 
+    def setup_openapi(self):
+        setup_aiohttp_apispec(
+            app=self,
+            title="Prompting API",
+            url="/docs/swagger.json",
+            swagger_path="/docs",
+        )
+
 
 def main(run_aio_app=True, test=False) -> None:
     loop = asyncio.get_event_loop()
@@ -48,7 +79,7 @@ def main(run_aio_app=True, test=False) -> None:
         # Instantiate the application with the actual validator
         bt.logging.info("Starting validator application.")
         validator_app = ValidatorApplication()
-        bt.logging.success(f"Validator app initialized successfully", validator_app)
+        bt.logging.success("Validator app initialized successfully", validator_app)
 
         try:
             web.run_app(validator_app, port=port, loop=loop)