Merge pull request #3 from antoninoLorenzo/docker_py

Implemented `ai_ops_api.py`

Dockerfile

@@ -1,6 +1,5 @@
 # ----------- Agent API Docker File
 # TODO : Setup volume for persistent sessions
-# TODO : add healthcheck
 
 # Kali Setup
 FROM kalilinux/kali-rolling
@@ -32,14 +31,9 @@ RUN cd AI-OPS/  && \
 
 VOLUME ["/root/.aiops"]
 
-# Run healtcheck
-# ...
-
 # Run API
 ENV MODEL=${ollama_model}
 ENV ENDPOINT=${ollama_endpoint}
-RUN echo "MODEL=$MODEL"       && \
-    echo "ENDPOINT=$ENDPOINT"
 
 EXPOSE 8000
 CMD ["fastapi", "dev", "--host", "0.0.0.0", "./AI-OPS/src/api.py"]

README.md

@@ -1,4 +1,4 @@
-![pylint](https://img.shields.io/badge/PyLint-8.44-yellow?logo=python&logoColor=white)
+![pylint](https://img.shields.io/badge/PyLint-8.58-yellow?logo=python&logoColor=white)
 
 🚧 *Under Development* 🚧
 
@@ -46,8 +46,9 @@ Issues:
 -->
 # 💻 Install
 **Requirements**
-- Ollama (see [Ollama](https://github.com/ollama/ollama))
-- Docker (see [Docker Desktop](https://docs.docker.com/desktop/)) (*in development*)
+- Python
+- [Ollama](https://github.com/ollama/ollama)
+- Docker (see [Docker Desktop](https://docs.docker.com/desktop/)) 
 
 ### End-User
 1. **Setup**
@@ -62,12 +63,9 @@ Issues:
 - As an alternative see my solution: [Ollama on Colab](https://github.com/antoninoLorenzo/Ollama-on-Colab-with-ngrok)
 
 3. **Agent API**
-- Build Docker Image
+- Run `ai_ops_api.py` or build manually:
   ```
   docker build -t ai-ops:api-dev --build-arg ollama_endpoint=ENDPOINT ollama_model=MODEL .
-  ```
-- Run Docker Container
-  ```
   docker run -p 8000:8000 ai-ops:api-dev
   ```
 

ai_ops_api.py

@@ -0,0 +1,107 @@
+import io
+import sys
+import argparse
+from pathlib import Path
+
+import docker
+from docker.models.containers import Container
+from docker.models.images import Image
+from docker.errors import DockerException
+
+try:
+    CLIENT = docker.from_env()
+except DockerException:
+    print('[!] Docker daemon not running.')
+    sys.exit(-1)
+
+DK_FILE_PATH = Path(Path(__file__).parent / 'Dockerfile')
+BUILD_ARGS = {
+    'ollama_model': 'gemma:7b',
+    'ollama_endpoint': 'http://localhost:11434'
+}
+
+
+def run_docker_api(container: Container | None = None,
+                   build_image: Image | None = None):
+    print('### run_docker_api')
+    img = None
+    if container:
+        print('Container is provided')
+        if not container.status == 'running':
+            container.start()
+    elif not container and build_image:
+        print('Image is provided')
+        img = build_image
+    elif not container and not build_image:
+        print('Building Image')
+        with open(str(DK_FILE_PATH), 'r') as fp:
+            dockerfile = fp.read()
+
+        img, logs = CLIENT.images.build(
+            fileobj=io.BytesIO(dockerfile.encode('utf-8')),
+            tag='ai-ops:api-dev',
+            rm=True,
+            buildargs=BUILD_ARGS,
+            pull=True,
+        )
+
+        for log in logs:
+            if 'stream' in log:
+                print(log['stream'], end='')
+
+    if img:
+        print('Run Container')
+        _ = CLIENT.containers.run(
+            img,
+            name='ai-ops-api',
+            tty=True,
+            stdin_open=True,
+            detach=True,
+            ports={"8000/tcp": 8000}
+        )
+
+    print('Done')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        '--update',
+        default="0",
+        choices=["0", "1"],
+        required=False,
+        help='Update: if set to 1 image is rebuilt'
+    )
+
+    args = parser.parse_args(sys.argv[1:])
+
+    # image gets built if we don't want to update it
+    # (`update` is False ) and it is not found
+    _update = int(args.update)
+    _container = None
+    image = None
+    available_images = {image.tags[0]: image for image in CLIENT.images.list()}
+    if not _update and 'ai-ops:api-dev' in available_images.keys():
+        print('Found Image for AI-OPS')
+        image = available_images['ai-ops:api-dev']
+
+    # container is restarted if we don't want to update it
+    # (`update` is False ) and it is found, if not found run from image
+    available_containers = {ctr.attrs.get('Name'): ctr for ctr in CLIENT.containers.list(all=True)}
+    if not _update and '/ai-ops-api' in available_containers.keys():
+        print('Found Container for AI-OPS')
+        _container = available_containers['/ai-ops-api']
+
+    # delete existing container if it already exists
+    if _update and '/ai-ops-api' in available_containers.keys():
+        ctr: Container = available_containers['/ai-ops-api']
+        if ctr.status == 'running':
+            ctr.stop()
+        ctr.remove()
+
+    # delete existing image if it already exists
+    if _update and 'ai-ops:api-dev' in available_images.keys():
+        available_images['ai-ops:api-dev'].remove()
+
+    run_docker_api(container=_container, build_image=image)

requirements.txt

@@ -18,3 +18,4 @@ google-generativeai
 pydantic_settings
 httpx
 tool-parse
+docker

src/agent/agent.py

@@ -3,7 +3,7 @@
 import json
 from json import JSONDecodeError
 
-from tool_parse import ToolRegistry, NotRegisteredError
+from tool_parse import ToolRegistry
 
 from src.agent.knowledge import Store
 from src.agent.llm import LLM, AVAILABLE_PROVIDERS
@@ -121,7 +121,7 @@ def invoke_tools(self, tool_response):
                 )
                 call_stack.append(tool_meta)
                 results.append({'role': 'tool', 'content': str(res)})
-            except NotRegisteredError:
+            except Exception:
                 pass
 
         return results

src/api.py

@@ -28,7 +28,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
 from pydantic_settings import BaseSettings
-from tool_parse import ToolRegistry, NotRegisteredError
+from tool_parse import ToolRegistry
 
 from src import initialize_knowledge
 from src.agent import Agent