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

Commit 4ee938f

Browse files
aponcedeleonchaponcedeleonch
committed
Reformatted checkf for is_fim_request
1 parent 0172668 commit 4ee938f

11 files changed

Lines changed: 129 additions & 59 deletions

File tree

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
from litellm import ChatCompletionRequest
22

3+
from codegate.codegate_logging import setup_logging
34
from codegate.pipeline.base import PipelineContext, PipelineResponse, PipelineResult, PipelineStep
45

6+
logger = setup_logging()
7+
58

69
class SecretAnalyzer(PipelineStep):
710
"""Pipeline step that handles version information requests."""
@@ -34,8 +37,9 @@ async def process(
3437
# message_with_secrets = any(messages_contain_secretes)
3538

3639
# For the moment to test shortcutting just treat all messages as if they contain secrets
37-
message_with_secrets = True
40+
message_with_secrets = False
3841
if message_with_secrets:
42+
logger.info('Blocking message with secrets.')
3943
return PipelineResult(
4044
response=PipelineResponse(
4145
step_name=self.name,

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
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+
LiteLLM Shim is a wrapper around LiteLLM's API that allows us to use it with
11+
our own completion handler interface without exposing the underlying
12+
LiteLLM API.
13+
"""
14+
15+
async def execute_completion(
16+
self,
17+
request: ChatCompletionRequest,
18+
api_key: Optional[str],
19+
stream: bool = False,
20+
) -> Union[ModelResponse, AsyncIterator[ModelResponse]]:
21+
"""
22+
Execute the completion request with LiteLLM's API
23+
"""
24+
model_in_request = request['model']
25+
if not model_in_request.startswith('anthropic/'):
26+
request['model'] = f'anthropic/{model_in_request}'
27+
return await super().execute_completion(request, api_key, stream)

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

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,18 +4,24 @@
44
from fastapi import Header, HTTPException, Request
55

66
from codegate.providers.anthropic.adapter import AnthropicInputNormalizer, AnthropicOutputNormalizer
7-
from codegate.providers.base import BaseProvider
8-
from codegate.providers.litellmshim import LiteLLmShim, anthropic_stream_generator
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
910

1011

1112
class AnthropicProvider(BaseProvider):
12-
def __init__(self, pipeline_processor=None):
13-
completion_handler = LiteLLmShim(stream_generator=anthropic_stream_generator)
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)
1419
super().__init__(
1520
AnthropicInputNormalizer(),
1621
AnthropicOutputNormalizer(),
1722
completion_handler,
1823
pipeline_processor,
24+
fim_pipeline_processor
1925
)
2026

2127
@property
@@ -40,5 +46,6 @@ async def create_message(
4046
body = await request.body()
4147
data = json.loads(body)
4248

43-
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)
4451
return self._completion_handler.create_streaming_response(stream)

‎src/codegate/providers/base.py‎

Lines changed: 65 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
11
from abc import ABC, abstractmethod
22
from typing import Any, AsyncIterator, Callable, Dict, Optional, Union
33

4-
from fastapi import APIRouter
4+
from fastapi import APIRouter, Request
55
from litellm import ModelResponse
66
from litellm.types.llms.openai import ChatCompletionRequest
77

8+
from codegate.codegate_logging import setup_logging
89
from codegate.pipeline.base import PipelineResult, SequentialPipelineProcessor
910
from codegate.providers.completion.base import BaseCompletionHandler
1011
from codegate.providers.formatting.input_pipeline import PipelineResponseFormatter
1112
from codegate.providers.normalizer.base import ModelInputNormalizer, ModelOutputNormalizer
1213

14+
logger = setup_logging()
1315
StreamGenerator = Callable[[AsyncIterator[Any]], AsyncIterator[str]]
1416

1517

@@ -25,12 +27,14 @@ def __init__(
2527
output_normalizer: ModelOutputNormalizer,
2628
completion_handler: BaseCompletionHandler,
2729
pipeline_processor: Optional[SequentialPipelineProcessor] = None,
30+
fim_pipeline_processor: Optional[SequentialPipelineProcessor] = None,
2831
):
2932
self.router = APIRouter()
3033
self._completion_handler = completion_handler
3134
self._input_normalizer = input_normalizer
3235
self._output_normalizer = output_normalizer
3336
self._pipeline_processor = pipeline_processor
37+
self._fim_pipelin_processor = fim_pipeline_processor
3438

3539
self._pipeline_response_formatter = PipelineResponseFormatter(output_normalizer)
3640

@@ -48,22 +52,76 @@ def provider_route_name(self) -> str:
4852
async def _run_input_pipeline(
4953
self,
5054
normalized_request: ChatCompletionRequest,
55+
is_fim_request: bool
5156
) -> PipelineResult:
52-
if self._pipeline_processor is None:
57+
# Decide which pipeline processor to use
58+
if is_fim_request:
59+
pipeline_processor = self._fim_pipelin_processor
60+
logger.info('FIM pipeline selected for execution.')
61+
else:
62+
pipeline_processor = self._pipeline_processor
63+
logger.info('Chat completion pipeline selected for execution.')
64+
if pipeline_processor is None:
5365
return PipelineResult(request=normalized_request)
5466

55-
result = await self._pipeline_processor.process_request(normalized_request)
67+
result = await pipeline_processor.process_request(normalized_request)
5668

5769
# TODO(jakub): handle this by returning a message to the client
5870
if result.error_message:
5971
raise Exception(result.error_message)
6072

6173
return result
6274

75+
def _is_fim_request_url(self, request: Request) -> bool:
76+
"""
77+
Checks the request URL to determine if a request is FIM or chat completion.
78+
Used by: llama.cpp
79+
"""
80+
request_path = request.url.path
81+
# Evaluate first a larger substring.
82+
if request_path.endswith("/chat/completions"):
83+
return False
84+
85+
if request_path.endswith("/completions"):
86+
return True
87+
88+
return False
89+
90+
def _is_fim_request_body(self, data: Dict) -> bool:
91+
"""
92+
Determine from the raw incoming data if it's a FIM request.
93+
Used by: OpenAI and Anthropic
94+
"""
95+
messages = data.get('messages', [])
96+
if not messages:
97+
return False
98+
99+
first_message_content = messages[0].get('content')
100+
if first_message_content is None:
101+
return False
102+
103+
fim_stop_sequences = ['</COMPLETION>', '<COMPLETION>', '</QUERY>', '<QUERY>']
104+
if isinstance(first_message_content, str):
105+
msg_prompt = first_message_content
106+
elif isinstance(first_message_content, list):
107+
msg_prompt = first_message_content[0].get('text', '')
108+
else:
109+
logger.warning(f'Could not determine if message was FIM from data: {data}')
110+
return False
111+
return all([stop_sequence in msg_prompt for stop_sequence in fim_stop_sequences])
112+
113+
def _is_fim_request(self, request: Request, data: Dict) -> bool:
114+
"""
115+
Determin if the request is FIM by the URL or the data of the request.
116+
"""
117+
# Avoid more expensive inspection of body by just checking the URL.
118+
if self._is_fim_request_url(request):
119+
return True
120+
121+
return self._is_fim_request_body(data)
122+
63123
async def complete(
64-
self,
65-
data: Dict,
66-
api_key: Optional[str],
124+
self, data: Dict, api_key: Optional[str], is_fim_request: bool
67125
) -> Union[ModelResponse, AsyncIterator[ModelResponse]]:
68126
"""
69127
Main completion flow with pipeline integration
@@ -79,7 +137,7 @@ async def complete(
79137
normalized_request = self._input_normalizer.normalize(data)
80138
streaming = data.get("stream", False)
81139

82-
input_pipeline_result = await self._run_input_pipeline(normalized_request)
140+
input_pipeline_result = await self._run_input_pipeline(normalized_request, is_fim_request)
83141
if input_pipeline_result.response:
84142
return self._pipeline_response_formatter.handle_pipeline_response(
85143
input_pipeline_result.response, streaming

‎src/codegate/providers/litellmshim/litellmshim.py‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,3 @@ def create_streaming_response(self, stream: AsyncIterator[Any]) -> StreamingResp
4343
},
4444
status_code=200,
4545
)
46-
47-
def is_fim_request(self, data: Dict) -> bool:
48-
"""
49-
Determine from the raw incoming data if it's a FIM request.
50-
This is needed here since completion_handler is used by provider and provider
51-
doesn't know about the adapter.
52-
"""
53-
return self._adapter.is_fim_request(data)

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -65,11 +65,3 @@ def create_streaming_response(self, stream: Iterator[Any]) -> StreamingResponse:
6565
},
6666
status_code=200,
6767
)
68-
69-
def is_fim_request(self, data: Dict) -> bool:
70-
"""
71-
Determine from the raw incoming data if it's a FIM request.
72-
This is needed here since completion_handler is used by provider and provider
73-
doesn't know about the adapter.
74-
"""
75-
return self._adapter.is_fim_request(data)

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

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,21 +3,24 @@
33

44
from fastapi import Request
55

6-
from codegate.pipeline.base import SequentialPipelineProcessor
7-
from codegate.providers.base import BaseProvider
8-
from codegate.providers.llamacpp.adapter import LlamaCppAdapter
6+
from codegate.providers.base import BaseProvider, SequentialPipelineProcessor
97
from codegate.providers.llamacpp.completion_handler import LlamaCppCompletionHandler
108
from codegate.providers.llamacpp.normalizer import LLamaCppInputNormalizer, LLamaCppOutputNormalizer
119

1210

1311
class LlamaCppProvider(BaseProvider):
14-
def __init__(self, pipeline_processor=None):
12+
def __init__(
13+
self,
14+
pipeline_processor: Optional[SequentialPipelineProcessor] = None,
15+
fim_pipeline_processor: Optional[SequentialPipelineProcessor] = None
16+
):
1517
completion_handler = LlamaCppCompletionHandler()
1618
super().__init__(
1719
LLamaCppInputNormalizer(),
1820
LLamaCppOutputNormalizer(),
1921
completion_handler,
2022
pipeline_processor,
23+
fim_pipeline_processor
2124
)
2225

2326
@property
@@ -37,5 +40,6 @@ async def create_completion(
3740
body = await request.body()
3841
data = json.loads(body)
3942

40-
stream = await self.complete(data, api_key=None)
43+
is_fim_request = self._is_fim_request(request, data)
44+
stream = await self.complete(data, None, is_fim_request=is_fim_request)
4145
return self._completion_handler.create_streaming_response(stream)

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

Lines changed: 10 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,24 @@
33

44
from fastapi import Header, HTTPException, Request
55

6-
from codegate.pipeline.base import SequentialPipelineProcessor
7-
from codegate.providers.base import BaseProvider
6+
from codegate.providers.base import BaseProvider, SequentialPipelineProcessor
87
from codegate.providers.litellmshim import LiteLLmShim, sse_stream_generator
98
from codegate.providers.openai.adapter import OpenAIInputNormalizer, OpenAIOutputNormalizer
109

1110

1211
class OpenAIProvider(BaseProvider):
13-
def __init__(self, pipeline_processor=None):
12+
def __init__(
13+
self,
14+
pipeline_processor: Optional[SequentialPipelineProcessor] = None,
15+
fim_pipeline_processor: Optional[SequentialPipelineProcessor] = None
16+
):
1417
completion_handler = LiteLLmShim(stream_generator=sse_stream_generator)
1518
super().__init__(
1619
OpenAIInputNormalizer(),
1720
OpenAIOutputNormalizer(),
1821
completion_handler,
1922
pipeline_processor,
23+
fim_pipeline_processor
2024
)
2125

2226
@property
@@ -31,22 +35,8 @@ def _setup_routes(self):
3135
"""
3236

3337
@self.router.post(f"/{self.provider_route_name}/chat/completions")
34-
async def create_completion(
35-
request: Request,
36-
authorization: str = Header(..., description="Bearer token"),
37-
):
38-
if not authorization.startswith("Bearer "):
39-
raise HTTPException(status_code=401, detail="Invalid authorization header")
40-
41-
api_key = authorization.split(" ")[1]
42-
body = await request.body()
43-
data = json.loads(body)
44-
45-
stream = await self.complete(data, api_key)
46-
return self._completion_handler.create_streaming_response(stream)
47-
4838
@self.router.post(f"/{self.provider_route_name}/completions")
49-
async def create_fim(
39+
async def create_completion(
5040
request: Request,
5141
authorization: str = Header(..., description="Bearer token"),
5242
):
@@ -57,5 +47,6 @@ async def create_fim(
5747
body = await request.body()
5848
data = json.loads(body)
5949

60-
stream = await self.complete(data, api_key)
50+
is_fim_request = self._is_fim_request(request, data)
51+
stream = await self.complete(data, api_key, is_fim_request=is_fim_request)
6152
return self._completion_handler.create_streaming_response(stream)

‎src/codegate/server.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44

55
from codegate import __description__, __version__
66
from codegate.pipeline.base import PipelineStep, SequentialPipelineProcessor
7-
from codegate.pipeline.version.version import CodegateVersion
87
from codegate.pipeline.fim.secret_analyzer import SecretAnalyzer
98
from codegate.pipeline.version.version import CodegateVersion
109
from codegate.providers.anthropic.provider import AnthropicProvider
@@ -26,7 +25,6 @@ def init_app() -> FastAPI:
2625
fim_steps: List[PipelineStep] = [
2726
SecretAnalyzer(),
2827
]
29-
3028
pipeline = SequentialPipelineProcessor(steps)
3129
fim_pipeline = SequentialPipelineProcessor(fim_steps)
3230

‎tests/providers/anthropic/test_adapter.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ def test_normalize_anthropic_input(input_normalizer):
4040
{"content": "You are an expert code reviewer", "role": "system"},
4141
{"content": [{"text": "Review this code", "type": "text"}], "role": "user"},
4242
],
43-
"model": "anthropic/claude-3-haiku-20240307",
43+
"model": "claude-3-haiku-20240307",
4444
"stream": True,
4545
}
4646

0 commit comments

Comments
 (0)