This repository was archived by the owner on Jun 5, 2025. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 91
Expand file tree
/
Copy pathfactory.py
More file actions
76 lines (69 loc) · 2.83 KB
/
factory.py
File metadata and controls
76 lines (69 loc) · 2.83 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
from typing import List
from codegate.clients.clients import ClientType
from codegate.config import Config
from codegate.pipeline.base import PipelineStep, SequentialPipelineProcessor
from codegate.pipeline.cli.cli import CodegateCli
from codegate.pipeline.codegate_context_retriever.codegate import CodegateContextRetriever
from codegate.pipeline.comment.output import CodeCommentStep
from codegate.pipeline.output import OutputPipelineProcessor, OutputPipelineStep
from codegate.pipeline.pii.pii import (
CodegatePii,
PiiRedactionNotifier,
PiiUnRedactionStep,
)
from codegate.pipeline.secrets.secrets import (
CodegateSecrets,
SecretRedactionNotifier,
SecretUnredactionStep,
)
from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
from codegate.pipeline.system_prompt.codegate import SystemPrompt
class PipelineFactory:
def __init__(self, sensitive_data_manager: SensitiveDataManager):
self.sensitive_data_manager = sensitive_data_manager
def create_input_pipeline(self, client_type: ClientType) -> SequentialPipelineProcessor:
input_steps: List[PipelineStep] = [
# make sure that this step is always first in the pipeline
# the other steps might send the request to a LLM for it to be analyzed
# and without obfuscating the secrets, we'd leak the secrets during those
# later steps
CodegateSecrets(),
CodegatePii(self.sensitive_data_manager),
CodegateCli(),
CodegateContextRetriever(),
SystemPrompt(
Config.get_config().prompts.default_chat, Config.get_config().prompts.client_prompts
),
]
return SequentialPipelineProcessor(
input_steps,
self.sensitive_data_manager,
client_type,
is_fim=False,
)
def create_fim_pipeline(self, client_type: ClientType) -> SequentialPipelineProcessor:
fim_steps: List[PipelineStep] = [
CodegateSecrets(),
CodegatePii(self.sensitive_data_manager),
]
return SequentialPipelineProcessor(
fim_steps,
self.sensitive_data_manager,
client_type,
is_fim=True,
)
def create_output_pipeline(self) -> OutputPipelineProcessor:
output_steps: List[OutputPipelineStep] = [
SecretRedactionNotifier(),
SecretUnredactionStep(),
PiiRedactionNotifier(),
PiiUnRedactionStep(),
CodeCommentStep(),
]
return OutputPipelineProcessor(output_steps)
def create_fim_output_pipeline(self) -> OutputPipelineProcessor:
fim_output_steps: List[OutputPipelineStep] = [
# temporarily disabled
# SecretUnredactionStep(),
]
return OutputPipelineProcessor(fim_output_steps)