Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Commit 6e7e646

Browse files
lukehindslukehinds
committed
Merge remote-tracking branch 'origin/main' into on-the-fly
2 parents 1090d43 + 189aee9 commit 6e7e646

28 files changed

Lines changed: 829 additions & 281 deletions

‎Dockerfile‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
# Builder stage: Install dependencies and build the application
2+
FROM python:3.12-slim AS builder
3+
4+
# Install system dependencies
5+
RUN apt-get update && apt-get install -y --no-install-recommends \
6+
gcc \
7+
g++ \
8+
&& rm -rf /var/lib/apt/lists/*
9+
10+
# Set environment variable to ensure Python modules are installed in the correct location
11+
ENV PYTHONPATH=/app
12+
13+
# Install Poetry
14+
RUN pip install poetry==1.8.4
15+
16+
# Create a non-root user and switch to it
17+
RUN adduser --system --no-create-home codegate --uid 1000
18+
19+
# Set the working directory
20+
WORKDIR /app
21+
22+
# Copy only the files needed for installing dependencies
23+
COPY pyproject.toml poetry.lock* /app/
24+
25+
# Configure Poetry and install dependencies
26+
RUN poetry config virtualenvs.create false && \
27+
poetry install --no-dev
28+
29+
# Copy the rest of the application
30+
COPY . /app
31+
32+
# Runtime stage: Create the final lightweight image
33+
FROM python:3.12-slim AS runtime
34+
35+
# Install runtime system dependencies
36+
RUN apt-get update && apt-get install -y --no-install-recommends \
37+
libgomp1 \
38+
&& rm -rf /var/lib/apt/lists/*
39+
40+
# Create a non-root user and switch to it
41+
RUN adduser --system --no-create-home codegate --uid 1000
42+
USER codegate
43+
44+
# Copy necessary artifacts from the builder stage
45+
COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages
46+
COPY --from=builder /app /app
47+
48+
# Set the working directory
49+
WORKDIR /app
50+
51+
# Set the PYTHONPATH environment variable
52+
ENV PYTHONPATH=/app/src
53+
54+
# Allow to expose weaviate_data volume
55+
VOLUME ["/app/weaviate_data"]
56+
57+
# Set the container's default entrypoint
58+
EXPOSE 8989
59+
#ENTRYPOINT ["python", "-m", "src.codegate.cli", "serve", "--port", "8989", "--host", "0.0.0.0"]
60+
CMD ["python", "-m", "src.codegate.cli", "serve", "--port", "8989", "--host", "0.0.0.0"]

‎Makefile‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
.PHONY: clean install format lint test security build all
2+
CONTAINER_BUILD?=docker buildx build
3+
VER?=0.1.0
24

35
clean:
46
rm -rf build/
@@ -27,4 +29,8 @@ security:
2729
build: clean test
2830
poetry build
2931

32+
image-build:
33+
$(CONTAINER_BUILD) -f Dockerfile -t codegate . -t ghcr.io/stacklok/codegate:$(VER) --load
34+
35+
3036
all: clean install format lint test security build

‎README.md‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,3 +108,19 @@ pip install -e ".[dev]"
108108
```bash
109109
pytest
110110
```
111+
112+
113+
### Running from image
114+
115+
A docker image can be built just with `make image-build`. That will start a codegate server ready to use.
116+
Then it can be started with:
117+
118+
```bash
119+
docker run -p 8989:8989 codegate:latest
120+
```
121+
122+
Additionally if you want to start with a pre-created database, a volume can be mounted:
123+
124+
```bash
125+
docker run -p 8989:8989 -v /path/to/volume:/app/weaviate_data codegate:latest
126+
```

‎src/codegate/pipeline/fim/secret_analyzer.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
from litellm import ChatCompletionRequest
2+
3+
from codegate.codegate_logging import setup_logging
4+
from codegate.pipeline.base import PipelineContext, PipelineResponse, PipelineResult, PipelineStep
5+
6+
logger = setup_logging()
7+
8+
9+
class SecretAnalyzer(PipelineStep):
10+
"""Pipeline step that handles analyzing secrets in FIM pipeline."""
11+
12+
message_blocked = """
13+
⚠️ CodeGate Security Warning! Analysis Report ⚠️
14+
Potential leak of sensitive credentials blocked
15+
16+
Recommendations:
17+
- Use environment variables for secrets
18+
"""
19+
20+
@property
21+
def name(self) -> str:
22+
"""
23+
Returns the name of this pipeline step.
24+
25+
Returns:
26+
str: The identifier 'fim-secret-analyzer'
27+
"""
28+
return "fim-secret-analyzer"
29+
30+
async def process(
31+
self,
32+
request: ChatCompletionRequest,
33+
context: PipelineContext
34+
) -> PipelineResult:
35+
# We should call here Secrets Blocking module to see if the request messages contain secrets
36+
# messages_contain_secrets = [analyze_msg_secrets(msg) for msg in request.messages]
37+
# message_with_secrets = any(messages_contain_secretes)
38+
39+
# For the moment to test shortcutting just treat all messages as if they contain secrets
40+
message_with_secrets = False
41+
if message_with_secrets:
42+
logger.info('Blocking message with secrets.')
43+
return PipelineResult(
44+
response=PipelineResponse(
45+
step_name=self.name,
46+
content=self.message_blocked,
47+
model=request["model"],
48+
),
49+
)
50+
51+
# No messages with secrets, execute the rest of the pipeline
52+
return PipelineResult(request=request)

‎src/codegate/providers/anthropic/adapter.py‎

Lines changed: 15 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,40 +1,30 @@
1-
from typing import Any, Dict, Optional
2-
3-
from litellm import AdapterCompletionStreamWrapper, ChatCompletionRequest, ModelResponse
41
from litellm.adapters.anthropic_adapter import (
52
AnthropicAdapter as LitellmAnthropicAdapter,
63
)
7-
from litellm.types.llms.anthropic import AnthropicResponse
84

9-
from codegate.providers.base import StreamGenerator
10-
from codegate.providers.litellmshim import BaseAdapter, anthropic_stream_generator
5+
from codegate.providers.litellmshim.adapter import (
6+
LiteLLMAdapterInputNormalizer,
7+
LiteLLMAdapterOutputNormalizer,
8+
)
119

1210

13-
class AnthropicAdapter(BaseAdapter):
11+
class AnthropicInputNormalizer(LiteLLMAdapterInputNormalizer):
1412
"""
1513
LiteLLM's adapter class interface is used to translate between the Anthropic data
1614
format and the underlying model. The AnthropicAdapter class contains the actual
1715
implementation of the interface methods, we just forward the calls to it.
1816
"""
1917

20-
def __init__(self, stream_generator: StreamGenerator = anthropic_stream_generator):
21-
self.litellm_anthropic_adapter = LitellmAnthropicAdapter()
22-
super().__init__(stream_generator)
18+
def __init__(self):
19+
super().__init__(LitellmAnthropicAdapter())
2320

24-
def translate_completion_input_params(
25-
self,
26-
completion_request: Dict,
27-
) -> Optional[ChatCompletionRequest]:
28-
return self.litellm_anthropic_adapter.translate_completion_input_params(completion_request)
2921

30-
def translate_completion_output_params(
31-
self, response: ModelResponse
32-
) -> Optional[AnthropicResponse]:
33-
return self.litellm_anthropic_adapter.translate_completion_output_params(response)
22+
class AnthropicOutputNormalizer(LiteLLMAdapterOutputNormalizer):
23+
"""
24+
LiteLLM's adapter class interface is used to translate between the Anthropic data
25+
format and the underlying model. The AnthropicAdapter class contains the actual
26+
implementation of the interface methods, we just forward the calls to it.
27+
"""
3428

35-
def translate_completion_output_params_streaming(
36-
self, completion_stream: Any
37-
) -> AdapterCompletionStreamWrapper | None:
38-
return self.litellm_anthropic_adapter.translate_completion_output_params_streaming(
39-
completion_stream
40-
)
29+
def __init__(self):
30+
super().__init__(LitellmAnthropicAdapter())

‎src/codegate/providers/anthropic/completion_handler.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
from typing import AsyncIterator, Optional, Union
2+
3+
from litellm import ChatCompletionRequest, ModelResponse
4+
5+
from codegate.providers.litellmshim import LiteLLmShim
6+
7+
8+
class AnthropicCompletion(LiteLLmShim):
9+
"""
10+
AnthropicCompletion used by the Anthropic provider to execute completions
11+
"""
12+
13+
async def execute_completion(
14+
self,
15+
request: ChatCompletionRequest,
16+
api_key: Optional[str],
17+
stream: bool = False,
18+
) -> Union[ModelResponse, AsyncIterator[ModelResponse]]:
19+
"""
20+
Ensures the model name is prefixed with 'anthropic/' to explicitly route to Anthropic's API.
21+
22+
LiteLLM automatically maps most model names, but prepending 'anthropic/' forces the request
23+
to Anthropic. This avoids issues with unrecognized names like 'claude-3-5-sonnet-latest',
24+
which LiteLLM doesn't accept as a valid Anthropic model. This safeguard may be unnecessary
25+
but ensures compatibility.
26+
27+
For more details, refer to the
28+
[LiteLLM Documentation](https://docs.litellm.ai/docs/providers/anthropic).
29+
"""
30+
model_in_request = request['model']
31+
if not model_in_request.startswith('anthropic/'):
32+
request['model'] = f'anthropic/{model_in_request}'
33+
return await super().execute_completion(request, api_key, stream)

‎src/codegate/providers/anthropic/provider.py‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,28 @@
11
import json
2+
from typing import Optional
23

34
from fastapi import Header, HTTPException, Request
45

5-
from codegate.providers.anthropic.adapter import AnthropicAdapter
6-
from codegate.providers.base import BaseProvider
7-
from codegate.providers.litellmshim import LiteLLmShim
6+
from codegate.providers.anthropic.adapter import AnthropicInputNormalizer, AnthropicOutputNormalizer
7+
from codegate.providers.anthropic.completion_handler import AnthropicCompletion
8+
from codegate.providers.base import BaseProvider, SequentialPipelineProcessor
9+
from codegate.providers.litellmshim import anthropic_stream_generator
810

911

1012
class AnthropicProvider(BaseProvider):
11-
def __init__(self, pipeline_processor=None):
12-
adapter = AnthropicAdapter()
13-
completion_handler = LiteLLmShim(adapter)
14-
super().__init__(completion_handler, pipeline_processor)
13+
def __init__(
14+
self,
15+
pipeline_processor: Optional[SequentialPipelineProcessor] = None,
16+
fim_pipeline_processor: Optional[SequentialPipelineProcessor] = None
17+
):
18+
completion_handler = AnthropicCompletion(stream_generator=anthropic_stream_generator)
19+
super().__init__(
20+
AnthropicInputNormalizer(),
21+
AnthropicOutputNormalizer(),
22+
completion_handler,
23+
pipeline_processor,
24+
fim_pipeline_processor
25+
)
1526

1627
@property
1728
def provider_route_name(self) -> str:
@@ -35,5 +46,6 @@ async def create_message(
3546
body = await request.body()
3647
data = json.loads(body)
3748

38-
stream = await self.complete(data, x_api_key)
49+
is_fim_request = self._is_fim_request(request, data)
50+
stream = await self.complete(data, x_api_key, is_fim_request)
3951
return self._completion_handler.create_streaming_response(stream)

0 commit comments

Comments
 (0)