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 90
Expand file tree
/
Copy pathmodels.py
More file actions
111 lines (92 loc) · 3.42 KB
/
models.py
File metadata and controls
111 lines (92 loc) · 3.42 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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
from enum import Enum
from typing import Any, Optional, Self
import pydantic
from codegate.clients.clients import ClientType
from codegate.db.models import MuxRule as DBMuxRule
from codegate.db.models import ProviderEndpoint as DBProviderEndpoint
from codegate.db.models import ProviderType
class MuxMatcherType(str, Enum):
"""
Represents the different types of matchers we support.
The 3 rules present match filenames and request types. They're used in conjunction with the
matcher field in the MuxRule model.
E.g.
- catch_all-> Always match
- filename_match and match: requests.py -> Match the request if the filename is requests.py
- fim_filename and match: main.py -> Match the request if the request type is fim
and the filename is main.py
NOTE: Removing or updating fields from this enum will require a migration.
Adding new fields is safe.
"""
# Always match this prompt
catch_all = "catch_all"
# Match based on the filename. It will match if there is a filename
# in the request that matches the matcher either extension or full name (*.py or main.py)
filename_match = "filename_match"
# Match based on fim request type. It will match if the request type is fim
fim_filename = "fim_filename"
# Match based on chat request type. It will match if the request type is chat
chat_filename = "chat_filename"
class MuxRule(pydantic.BaseModel):
"""
Represents a mux rule for a provider.
"""
provider_name: str
provider_type: ProviderType
model: str
# The type of matcher to use
matcher_type: MuxMatcherType
# The actual matcher to use. Note that
# this depends on the matcher type.
matcher: Optional[str] = None
@classmethod
def from_db_models(
cls, db_mux_rule: DBMuxRule, db_provider_endpoint: DBProviderEndpoint
) -> Self:
"""
Convert a DBMuxRule and DBProviderEndpoint to a MuxRule.
"""
return cls(
provider_name=db_provider_endpoint.name,
provider_type=db_provider_endpoint.provider_type,
model=db_mux_rule.provider_model_name,
matcher_type=MuxMatcherType(db_mux_rule.matcher_type),
matcher=db_mux_rule.matcher_blob,
)
@classmethod
def from_mux_rule_with_provider_id(cls, rule: "MuxRuleWithProviderId") -> Self:
"""
Convert a MuxRuleWithProviderId to a MuxRule.
"""
return cls(
provider_name=rule.provider_name,
provider_type=rule.provider_type,
model=rule.model,
matcher_type=rule.matcher_type,
matcher=rule.matcher,
)
class MuxRuleWithProviderId(MuxRule):
"""
Represents a mux rule for a provider with provider ID.
Used internally for referring to a mux rule.
"""
provider_id: str
@classmethod
def from_db_models(
cls, db_mux_rule: DBMuxRule, db_provider_endpoint: DBProviderEndpoint
) -> Self:
"""
Convert a DBMuxRule and DBProviderEndpoint to a MuxRuleWithProviderId.
"""
return cls(
**MuxRule.from_db_models(db_mux_rule, db_provider_endpoint).model_dump(),
provider_id=db_mux_rule.provider_endpoint_id,
)
class ThingToMatchMux(pydantic.BaseModel):
"""
Represents the fields we can use to match a mux rule.
"""
body: Any
url_request_path: str
is_fim_request: bool
client_type: ClientType