Merge branch 'main' of github.com:stacklok/codegate into feat/555/version-endpoint

Author: alex-mcgovern

.github/workflows/image-publish.yml

@@ -87,6 +87,7 @@ jobs:
           cache-to: type=gha,mode=max
           build-args: |
             LATEST_RELEASE=${{ env.LATEST_RELEASE }}
+            CODEGATE_VERSION=${{ steps.version-string.outputs.tag }}
       - name: Capture Image Digest
         id: image-digest
         run: |

Dockerfile

@@ -1,6 +1,8 @@
 # Builder stage: Install dependencies and build the application
 FROM python:3.12-slim AS builder
 
+ARG CODEGATE_VERSION=dev
+
 # Install system dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \
     gcc \
@@ -21,6 +23,9 @@ RUN poetry config virtualenvs.create false && \
 # Copy the rest of the application
 COPY . /app
 
+# Overwrite the _VERSION variable in the code
+RUN sed -i "s/_VERSION =.*/_VERSION = \"${CODEGATE_VERSION}\"/g" /app/src/codegate/__init__.py
+
 # Build the webapp
 FROM node:23-slim AS webbuilder
 

Makefile

@@ -34,6 +34,7 @@ image-build:
 	DOCKER_BUILDKIT=1 $(CONTAINER_BUILD) \
 		-f Dockerfile \
 		--build-arg LATEST_RELEASE=$(curl -s "https://api.github.com/repos/stacklok/codegate-ui/releases/latest" | grep '"zipball_url":' | cut -d '"' -f 4) \
+		--build-arg CODEGATE_VERSION="$(shell git describe --tags --abbrev=0)-$(shell git rev-parse --short HEAD)-dev" \
 		-t codegate \
 		. \
 		-t ghcr.io/stacklok/codegate:$(VER) \

signatures.yaml

@@ -1,7 +1,6 @@
 ---
 - Amazon:
     - Access Key: (?:A3T[A-Z0-9]|AKIA|AGPA|AIDA|AROA|AIPA|ANPA|ANVA|ASIA|ABIA|ACCA)[A-Z0-9]{16}
-    - Secret Access Key: (?<![A-Za-z0-9/+])[A-Za-z0-9+=][A-Za-z0-9/+=]{38}[A-Za-z0-9+=](?![A-Za-z0-9/+=])
     # - Cognito User Pool ID: (?i)us-[a-z]{2,}-[a-z]{4,}-\d{1,}
     - RDS Password: (?i)(rds\-master\-password|db\-password)
     - SNS Confirmation URL: (?i)https:\/\/sns\.[a-z0-9-]+\.amazonaws\.com\/?Action=ConfirmSubscription&Token=[a-zA-Z0-9-=_]+

src/codegate/__init__.py

@@ -7,12 +7,19 @@
 from codegate.config import Config
 from codegate.exceptions import ConfigurationError
 
-try:
-    __version__ = metadata.version("codegate")
-    __description__ = metadata.metadata("codegate")["Summary"]
-except metadata.PackageNotFoundError:  # pragma: no cover
-    __version__ = "unknown"
-    __description__ = "codegate"
+_VERSION = "dev"
+_DESC = "CodeGate - A Generative AI security gateway."
+
+def __get_version_and_description() -> tuple[str, str]:
+    try:
+        version = metadata.version("codegate")
+        description = metadata.metadata("codegate")["Summary"]
+    except metadata.PackageNotFoundError:
+        version = _VERSION
+        description = _DESC
+    return version, description
+
+__version__, __description__ = __get_version_and_description()
 
 __all__ = ["Config", "ConfigurationError", "LogFormat", "LogLevel", "setup_logging"]
 

src/codegate/pipeline/base.py

@@ -239,36 +239,38 @@ def get_last_user_message(
     @staticmethod
     def get_last_user_message_block(
         request: ChatCompletionRequest,
-    ) -> Optional[str]:
+    ) -> Optional[tuple[str, int]]:
         """
         Get the last block of consecutive 'user' messages from the request.
 
         Args:
             request (ChatCompletionRequest): The chat completion request to process
 
         Returns:
-            Optional[str]: A string containing all consecutive user messages in the
+            Optional[str, int]: A string containing all consecutive user messages in the
                         last user message block, separated by newlines, or None if
                         no user message block is found.
+                        Index of the first message detected in the block.
         """
         if request.get("messages") is None:
             return None
 
         user_messages = []
         messages = request["messages"]
+        block_start_index = None
 
         # Iterate in reverse to find the last block of consecutive 'user' messages
         for i in reversed(range(len(messages))):
             if messages[i]["role"] == "user" or messages[i]["role"] == "assistant":
-                content_str = None
-                if "content" in messages[i]:
-                    content_str = messages[i]["content"]  # type: ignore
-                else:
+                content_str = messages[i].get("content")
+                if content_str is None:
                     continue
 
                 if messages[i]["role"] == "user":
                     user_messages.append(content_str)
-                # specifically for Aider, when "ok." block is found, stop
+                    block_start_index = i
+
+                # Specifically for Aider, when "Ok." block is found, stop
                 if content_str == "Ok." and messages[i]["role"] == "assistant":
                     break
             else:
@@ -277,8 +279,9 @@ def get_last_user_message_block(
                     break
 
         # Reverse the collected user messages to preserve the original order
-        if user_messages:
-            return "\n".join(reversed(user_messages))
+        if user_messages and block_start_index is not None:
+            content = "\n".join(reversed(user_messages))
+            return content, block_start_index
 
         return None
 

src/codegate/pipeline/codegate_context_retriever/codegate.py

@@ -60,9 +60,10 @@ async def process(
         Use RAG DB to add context to the user request
         """
         # Get the latest user message
-        user_message = self.get_last_user_message_block(request)
-        if not user_message:
+        last_message = self.get_last_user_message_block(request)
+        if not last_message:
             return PipelineResult(request=request)
+        user_message, _ = last_message
 
         # Create storage engine object
         storage_engine = StorageEngine()

src/codegate/pipeline/extract_snippets/extract_snippets.py

@@ -151,9 +151,10 @@ async def process(
         request: ChatCompletionRequest,
         context: PipelineContext,
     ) -> PipelineResult:
-        msg_content = self.get_last_user_message_block(request)
-        if not msg_content:
+        last_message = self.get_last_user_message_block(request)
+        if not last_message:
             return PipelineResult(request=request, context=context)
+        msg_content, _ = last_message
         snippets = extract_snippets(msg_content)
 
         logger.info(f"Extracted {len(snippets)} code snippets from the user message")

src/codegate/pipeline/secrets/secrets.py

@@ -271,11 +271,12 @@ async def process(
         new_request = request.copy()
         total_matches = []
 
-        # Process all messages
+        # get last user message block to get index for the first relevant user message
+        last_user_message = self.get_last_user_message_block(new_request)
         last_assistant_idx = -1
-        for i, message in enumerate(new_request["messages"]):
-            if message.get("role", "") == "assistant":
-                last_assistant_idx = i
+        if last_user_message:
+            _, user_idx = last_user_message
+            last_assistant_idx = user_idx - 1
 
         # Process all messages
         for i, message in enumerate(new_request["messages"]):
@@ -312,8 +313,8 @@ class SecretUnredactionStep(OutputPipelineStep):
     """Pipeline step that unredacts protected content in the stream"""
 
     def __init__(self):
-        self.redacted_pattern = re.compile(r"REDACTED<\$([^>]+)>")
-        self.marker_start = "REDACTED<$"
+        self.redacted_pattern = re.compile(r"REDACTED<(\$?[^>]+)>")
+        self.marker_start = "REDACTED<"
         self.marker_end = ">"
 
     @property
@@ -365,6 +366,8 @@ async def process_chunk(
         if match:
             # Found a complete marker, process it
             encrypted_value = match.group(1)
+            if encrypted_value.startswith('$'):
+                encrypted_value = encrypted_value[1:]
             original_value = input_context.sensitive.manager.get_original_value(
                 encrypted_value,
                 input_context.sensitive.session_id,
@@ -399,7 +402,7 @@ async def process_chunk(
             return []
 
         if self._is_partial_marker_prefix(buffered_content):
-            context.prefix_buffer += buffered_content
+            context.prefix_buffer = buffered_content
             return []
 
         # No markers or partial markers, let pipeline handle the chunk normally

src/codegate/providers/ollama/completion_handler.py

@@ -1,3 +1,4 @@
+import json
 from typing import AsyncIterator, Optional, Union
 
 import structlog
@@ -11,18 +12,21 @@
 
 
 async def ollama_stream_generator(
-    stream: AsyncIterator[ChatResponse],
+    stream: AsyncIterator[ChatResponse]
 ) -> AsyncIterator[str]:
     """OpenAI-style SSE format"""
     try:
         async for chunk in stream:
-            print(chunk)
             try:
-                yield f"{chunk.model_dump_json()}\n\n"
+                content = chunk.model_dump_json()
+                if content:
+                    yield f"{chunk.model_dump_json()}\n"
             except Exception as e:
-                yield f"{str(e)}\n\n"
+                if str(e):
+                    yield f"{str(e)}\n"
     except Exception as e:
-        yield f"{str(e)}\n\n"
+        if str(e):
+            yield f"{str(e)}\n"
 
 
 class OllamaShim(BaseCompletionHandler):
@@ -39,17 +43,17 @@ async def execute_completion(
     ) -> Union[ChatResponse, GenerateResponse]:
         """Stream response directly from Ollama API."""
         if is_fim_request:
-            prompt = request["messages"][0]["content"]
+            prompt = request["messages"][0].get("content", "")
             response = await self.client.generate(
-                model=request["model"], prompt=prompt, stream=stream, options=request["options"]
+                model=request["model"], prompt=prompt, stream=stream, options=request["options"]  # type: ignore
             )
         else:
             response = await self.client.chat(
                 model=request["model"],
                 messages=request["messages"],
-                stream=stream,
-                options=request["options"],
-            )
+                stream=stream,  # type: ignore
+                options=request["options"],  # type: ignore
+            )  # type: ignore
         return response
 
     def _create_streaming_response(self, stream: AsyncIterator[ChatResponse]) -> StreamingResponse: