Skip to content

Commit 623beb7

Browse files
ishworgiri1999ishworgiri1999
committed
rearrange api
1 parent a676707 commit 623beb7

26 files changed

+654
-493
lines changed

‎Dockerfile.queue

+1-1
Original file line numberDiff line numberDiff line change
@@ -13,4 +13,4 @@ RUN pip install -r requirements.txt
1313
# Make port 8000 available to the world outside this container
1414
EXPOSE 8000
1515

16-
CMD ["/bin/bash", "-c", "cd src && celery -A task.worker worker --loglevel=info --concurrency 1"]
16+
CMD ["/bin/bash", "-c", "cd src && celery -A worker.worker worker --loglevel=info --concurrency 1"]

‎Makefile

+1-1
Original file line numberDiff line numberDiff line change
@@ -27,4 +27,4 @@ dev:
2727

2828

2929
start-worker:
30-
cd src && celery -A task.worker worker --loglevel=info --concurrency 1
30+
cd src && celery -A worker.worker worker --loglevel=info --concurrency 1

‎src/ai/LLM/Strategies/AnthropicService.py

+51-30
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import os
21
from typing import Dict, List, Optional, Union
32

43
from anthropic import Anthropic
@@ -12,23 +11,31 @@
1211
SEARCH_TERMS_TEMPLATE,
1312
CONVERT_DICT_TO_STR_TEMPLATE,
1413
META_PROMPT_GENERATOR_TEMPLATE,
15-
LONG_RECOMMENDATION_TEMPLATE
14+
LONG_RECOMMENDATION_TEMPLATE,
1615
)
1716
from utils.text_tools import clean
18-
17+
from config import config
1918
import logging
19+
import json
20+
2021
logger = logging.getLogger(__name__)
2122

2223

2324
class AnthropicService(BaseLLMService):
24-
def __init__(self, api_key: str = os.getenv("ANTHROPIC_API_KEY", None), model: str = "claude-3-5-sonnet-20240620"):
25+
def __init__(
26+
self,
27+
api_key: str = config.anthropic_api_key,
28+
model: str = "claude-3-5-sonnet-20240620",
29+
):
2530
if api_key is None:
26-
raise ValueError("API key not provided and ANTHROPIC_API_KEY environment variable not set.")
31+
raise ValueError(
32+
"API key not provided and ANTHROPIC_API_KEY environment variable not set."
33+
)
2734
self.client = Anthropic(api_key=api_key)
2835
self.model = model
2936

3037
def get_model_name(self) -> str:
31-
return "-".join(self.model.split('-')[:-1])
38+
return "-".join(self.model.split("-")[:-1])
3239

3340
def get_url(self) -> str:
3441
return "-"
@@ -42,19 +49,23 @@ def generate(self, prompt: str) -> Dict[str, str]:
4249
content = message.content[0].text
4350
return {"response": content}
4451

45-
def classify_kind(self, finding: Finding, options: Optional[List[FindingKind]] = None) -> FindingKind:
52+
def classify_kind(
53+
self, finding: Finding, options: Optional[List[FindingKind]] = None
54+
) -> FindingKind:
4655
if options is None:
4756
options = list(FindingKind)
4857

49-
options_str = ', '.join([kind.name for kind in options])
58+
options_str = ", ".join([kind.name for kind in options])
5059
prompt = CLASSIFY_KIND_TEMPLATE.format(options=options_str, data=str(finding))
5160
response = self.generate(prompt)
52-
if 'response' not in response:
61+
if "response" not in response:
5362
logger.warning(f"Failed to classify the finding: {finding.title}")
5463
return FindingKind.DEFAULT
55-
return FindingKind[response['response'].strip()]
64+
return FindingKind[response["response"].strip()]
5665

57-
def get_recommendation(self, finding: Finding, short: bool = True) -> Union[str, List[str]]:
66+
def get_recommendation(
67+
self, finding: Finding, short: bool = True
68+
) -> Union[str, List[str]]:
5869
if short:
5970
prompt = SHORT_RECOMMENDATION_TEMPLATE.format(data=str(finding))
6071
else: # long recommendation
@@ -64,45 +75,55 @@ def get_recommendation(self, finding: Finding, short: bool = True) -> Union[str,
6475
else:
6576
prompt = GENERIC_LONG_RECOMMENDATION_TEMPLATE
6677

67-
finding.solution.add_to_metadata(f"prompt_{'short' if short else 'long'}", prompt)
78+
finding.solution.add_to_metadata(
79+
f"prompt_{'short' if short else 'long'}", prompt
80+
)
6881
response = self.generate(prompt)
6982

70-
if 'response' not in response:
83+
if "response" not in response:
7184
logger.warning(
72-
f"Failed to generate a {'short' if short else 'long'} recommendation for the finding: {finding.title}")
73-
return '' if short else ['']
85+
f"Failed to generate a {'short' if short else 'long'} recommendation for the finding: {finding.title}"
86+
)
87+
return "" if short else [""]
7488

75-
return clean(response['response'], llm_service=self)
89+
return clean(response["response"], llm_service=self)
7690

7791
def _generate_prompt_with_meta_prompts(self, finding: Finding) -> str:
7892
short_recommendation = finding.solution.short_description
79-
meta_prompt_generator = META_PROMPT_GENERATOR_TEMPLATE.format(finding=str(finding))
80-
meta_prompt_response = self.generate(meta_prompt_generator)
81-
meta_prompts = clean(
82-
meta_prompt_response.get("response", ""), llm_service=self
93+
meta_prompt_generator = META_PROMPT_GENERATOR_TEMPLATE.format(
94+
finding=str(finding)
8395
)
96+
meta_prompt_response = self.generate(meta_prompt_generator)
97+
meta_prompts = clean(meta_prompt_response.get("response", ""), llm_service=self)
8498

8599
long_prompt = LONG_RECOMMENDATION_TEMPLATE.format(meta_prompts=meta_prompts)
86100

87-
finding.solution.add_to_metadata("prompt_long_breakdown", {
88-
"short_recommendation": short_recommendation,
89-
"meta_prompts": meta_prompts
90-
})
101+
finding.solution.add_to_metadata(
102+
"prompt_long_breakdown",
103+
{
104+
"short_recommendation": short_recommendation,
105+
"meta_prompts": meta_prompts,
106+
},
107+
)
91108

92109
return long_prompt
93110

94111
def get_search_terms(self, finding: Finding) -> str:
95112
prompt = SEARCH_TERMS_TEMPLATE.format(data=str(finding))
96113
response = self.generate(prompt)
97-
if 'response' not in response:
98-
logger.warning(f"Failed to generate search terms for the finding: {finding.title}")
114+
if "response" not in response:
115+
logger.warning(
116+
f"Failed to generate search terms for the finding: {finding.title}"
117+
)
99118
return ""
100-
return clean(response['response'], llm_service=self)
119+
return clean(response["response"], llm_service=self)
101120

102121
def convert_dict_to_str(self, data) -> str:
103122
prompt = CONVERT_DICT_TO_STR_TEMPLATE.format(data=json.dumps(data))
104123
response = self.generate(prompt)
105-
if 'response' not in response:
106-
logger.info(f"Failed to convert dictionary to string, returning it as str conversion.")
124+
if "response" not in response:
125+
logger.info(
126+
f"Failed to convert dictionary to string, returning it as str conversion."
127+
)
107128
return str(data)
108-
return clean(response['response'], llm_service=self)
129+
return clean(response["response"], llm_service=self)

‎src/ai/LLM/Strategies/OLLAMAService.py

+14-15
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,11 @@
1919
CONVERT_DICT_TO_STR_TEMPLATE,
2020
)
2121

22+
from config import config
23+
2224
logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(name)s | %(message)s")
2325
logger = logging.getLogger(__name__)
2426

25-
import dotenv
26-
27-
dotenv.load_dotenv()
28-
2927

3028
def singleton(cls):
3129
instances = {}
@@ -41,8 +39,8 @@ def wrapper(*args, **kwargs):
4139
def is_up() -> bool:
4240
# res = requests.post(os.getenv('OLLAMA_URL') + '/api/show', json={'name': os.getenv('OLLAMA_MODEL', 'llama3')})
4341
res = httpx.post(
44-
os.getenv("OLLAMA_URL") + "/api/show",
45-
json={"name": os.getenv("OLLAMA_MODEL", "llama3")},
42+
config.ollama_url + "/api/show",
43+
json={"name": config.ollama_model},
4644
)
4745
if res.status_code == 200:
4846
return True
@@ -58,7 +56,7 @@ class OLLAMAService(BaseLLMService):
5856
"""
5957

6058
def __init__(
61-
self, model_url: Optional[str] = None, model_name: Optional[str] = None
59+
self, model_url: Optional[str] = None, model_name: Optional[str] = None
6260
):
6361
"""
6462
Initialize the LLMService object.
@@ -67,15 +65,12 @@ def __init__(
6765
"""
6866
# Configure logging level for httpx
6967
logging.getLogger("httpx").setLevel(logging.WARNING)
70-
print("model", os.getenv("OLLAMA_MODEL"))
7168
# Now, variables
7269
if model_url is None:
73-
model_url = os.getenv("OLLAMA_URL", "http://localhost:11434")
70+
model_url = config.ollama_url
7471
self.pull_url: str = model_url + "/api/pull"
7572
self.generate_url: str = model_url + "/api/generate"
76-
self.model_name: str = model_name or os.getenv(
77-
"OLLAMA_MODEL", "llama3:instruct"
78-
)
73+
self.model_name: str = model_name or config.ollama_model
7974

8075
self.generate_payload: Dict[str, Union[str, bool]] = {
8176
"model": self.model_name,
@@ -93,7 +88,9 @@ def init_pull_model(self) -> None:
9388
"""
9489
payload = {"name": self.model_name}
9590
try:
96-
response = httpx.post(self.pull_url, json=payload, timeout=60 * 10) # 10 minutes timeout
91+
response = httpx.post(
92+
self.pull_url, json=payload, timeout=60 * 10
93+
) # 10 minutes timeout
9794
response.raise_for_status()
9895
except httpx.ConnectError as e:
9996
logger.error(f"Failed to connect to the OLLAMA server: {e}")
@@ -133,14 +130,16 @@ def generate(self, prompt: str) -> Dict[str, str]:
133130
json_response = response.json()
134131
return parse_json(json_response["response"], strict=False)
135132
except json.JSONDecodeError as e:
136-
logger.error(f"LLM-Models JSON response is malformed, could not be parsed: {e}")
133+
logger.error(
134+
f"LLM-Models JSON response is malformed, could not be parsed: {e}"
135+
)
137136
return {}
138137
except httpx.ReadTimeout as e:
139138
logger.warning(f"ReadTimeout occurred: {e}")
140139
return {}
141140

142141
def classify_kind(
143-
self, finding: Finding, options: Optional[List[FindingKind]] = None
142+
self, finding: Finding, options: Optional[List[FindingKind]] = None
144143
) -> FindingKind:
145144
"""
146145
Classify the kind of security finding.

‎src/ai/LLM/Strategies/OpenAIService.py

+49-29
Original file line numberDiff line numberDiff line change
@@ -11,18 +11,25 @@
1111
SHORT_RECOMMENDATION_TEMPLATE,
1212
GENERIC_LONG_RECOMMENDATION_TEMPLATE,
1313
SEARCH_TERMS_TEMPLATE,
14-
CONVERT_DICT_TO_STR_TEMPLATE, META_PROMPT_GENERATOR_TEMPLATE, LONG_RECOMMENDATION_TEMPLATE
14+
CONVERT_DICT_TO_STR_TEMPLATE,
15+
META_PROMPT_GENERATOR_TEMPLATE,
16+
LONG_RECOMMENDATION_TEMPLATE,
1517
)
1618
from utils.text_tools import clean
1719

20+
from config import config
21+
1822
import logging
23+
1924
logger = logging.getLogger(__name__)
2025

2126

2227
class OpenAIService(BaseLLMService):
23-
def __init__(self, api_key: str = os.getenv("OPENAI_API_KEY", None), model: str = "gpt-4o"):
28+
def __init__(self, api_key: str = config.openai_api_key, model: str = "gpt-4o"):
2429
if api_key is None:
25-
raise ValueError("API key not provided and OPENAI_API_KEY environment variable not set.")
30+
raise ValueError(
31+
"API key not provided and OPENAI_API_KEY environment variable not set."
32+
)
2633
openai.api_key = api_key
2734
self.model = model
2835

@@ -34,25 +41,28 @@ def get_url(self) -> str:
3441

3542
def generate(self, prompt: str) -> Dict[str, str]:
3643
response = openai.chat.completions.create(
37-
model=self.model,
38-
messages=[{"role": "user", "content": prompt}]
44+
model=self.model, messages=[{"role": "user", "content": prompt}]
3945
)
4046
content = response.choices[0].message.content
4147
return {"response": content}
4248

43-
def classify_kind(self, finding: Finding, options: Optional[List[FindingKind]] = None) -> FindingKind:
49+
def classify_kind(
50+
self, finding: Finding, options: Optional[List[FindingKind]] = None
51+
) -> FindingKind:
4452
if options is None:
4553
options = list(FindingKind)
4654

47-
options_str = ', '.join([kind.name for kind in options])
55+
options_str = ", ".join([kind.name for kind in options])
4856
prompt = CLASSIFY_KIND_TEMPLATE.format(options=options_str, data=str(finding))
4957
response = self.generate(prompt)
50-
if 'response' not in response:
58+
if "response" not in response:
5159
logger.warning(f"Failed to classify the finding: {finding.title}")
5260
return FindingKind.DEFAULT
53-
return FindingKind[response['response'].strip()]
61+
return FindingKind[response["response"].strip()]
5462

55-
def get_recommendation(self, finding: Finding, short: bool = True) -> Union[str, List[str]]:
63+
def get_recommendation(
64+
self, finding: Finding, short: bool = True
65+
) -> Union[str, List[str]]:
5666
if short:
5767
prompt = SHORT_RECOMMENDATION_TEMPLATE.format(data=str(finding))
5868
else: # long recommendation
@@ -62,45 +72,55 @@ def get_recommendation(self, finding: Finding, short: bool = True) -> Union[str,
6272
else:
6373
prompt = GENERIC_LONG_RECOMMENDATION_TEMPLATE
6474

65-
finding.solution.add_to_metadata(f"prompt_{'short' if short else 'long'}", prompt)
75+
finding.solution.add_to_metadata(
76+
f"prompt_{'short' if short else 'long'}", prompt
77+
)
6678
response = self.generate(prompt)
6779

68-
if 'response' not in response:
80+
if "response" not in response:
6981
logger.warning(
70-
f"Failed to generate a {'short' if short else 'long'} recommendation for the finding: {finding.title}")
71-
return '' if short else ['']
82+
f"Failed to generate a {'short' if short else 'long'} recommendation for the finding: {finding.title}"
83+
)
84+
return "" if short else [""]
7285

73-
return clean(response['response'], llm_service=self)
86+
return clean(response["response"], llm_service=self)
7487

7588
def _generate_prompt_with_meta_prompts(self, finding: Finding) -> str:
7689
short_recommendation = finding.solution.short_description
77-
meta_prompt_generator = META_PROMPT_GENERATOR_TEMPLATE.format(finding=str(finding))
78-
meta_prompt_response = self.generate(meta_prompt_generator)
79-
meta_prompts = clean(
80-
meta_prompt_response.get("response", ""), llm_service=self
90+
meta_prompt_generator = META_PROMPT_GENERATOR_TEMPLATE.format(
91+
finding=str(finding)
8192
)
93+
meta_prompt_response = self.generate(meta_prompt_generator)
94+
meta_prompts = clean(meta_prompt_response.get("response", ""), llm_service=self)
8295

8396
long_prompt = LONG_RECOMMENDATION_TEMPLATE.format(meta_prompts=meta_prompts)
8497

85-
finding.solution.add_to_metadata("prompt_long_breakdown", {
86-
"short_recommendation": short_recommendation,
87-
"meta_prompts": meta_prompts
88-
})
98+
finding.solution.add_to_metadata(
99+
"prompt_long_breakdown",
100+
{
101+
"short_recommendation": short_recommendation,
102+
"meta_prompts": meta_prompts,
103+
},
104+
)
89105

90106
return long_prompt
91107

92108
def get_search_terms(self, finding: Finding) -> str:
93109
prompt = SEARCH_TERMS_TEMPLATE.format(data=str(finding))
94110
response = self.generate(prompt)
95-
if 'response' not in response:
96-
logger.warning(f"Failed to generate search terms for the finding: {finding.title}")
111+
if "response" not in response:
112+
logger.warning(
113+
f"Failed to generate search terms for the finding: {finding.title}"
114+
)
97115
return ""
98-
return clean(response['response'], llm_service=self)
116+
return clean(response["response"], llm_service=self)
99117

100118
def convert_dict_to_str(self, data) -> str:
101119
prompt = CONVERT_DICT_TO_STR_TEMPLATE.format(data=json.dumps(data))
102120
response = self.generate(prompt)
103-
if 'response' not in response:
104-
logger.info(f"Failed to convert dictionary to string, returning it as str conversion.")
121+
if "response" not in response:
122+
logger.info(
123+
f"Failed to convert dictionary to string, returning it as str conversion."
124+
)
105125
return str(data)
106-
return clean(response['response'], llm_service=self)
126+
return clean(response["response"], llm_service=self)

0 commit comments

Comments
 (0)