fixed gpt, added ollama

Author: DanMcInerney

Diff for: .env.example

@@ -4,7 +4,7 @@
 
 #ANTHROPIC_API_KEY=sk-1234
 ANTHROPIC_BASE_URL=https://api.anthropic.com
-ANTHROPIC_MODEL=claude-3-5-sonnet-20240620
+ANTHROPIC_MODEL=claude-3-5-sonnet-latest
 
 
 # For usage with OpenAI (-l gpt)
@@ -13,7 +13,7 @@ ANTHROPIC_MODEL=claude-3-5-sonnet-20240620
 
 #OPENAI_API_KEY=sk-1234
 OPENAI_BASE_URL=https://api.openai.com/v1
-OPENAI_MODEL=gpt-4o-2024-08-06
+OPENAI_MODEL=chatgpt-4o-latest
 
 
 # For usage with Ollama (-l ollama)

Diff for: vulnhuntr/LLMs.py

@@ -42,7 +42,7 @@ def _validate_response(self, response_text: str, response_model: BaseModel) -> B
                 response_text = self.prefill + response_text
             return response_model.model_validate_json(response_text)
         except ValidationError as e:
-            log.warning("Response validation failed", exc_info=e)
+            log.warning("[-] Response validation failed\n", exc_info=e)
             raise LLMError("Validation failed") from e
             # try:
             #     response_clean_attempt = response_text.split('{', 1)[1]
@@ -75,10 +75,10 @@ def chat(self, user_prompt: str, response_model: BaseModel = None, max_tokens: i
         return response_text
 
 class Claude(LLM):
-    def __init__(self, model: str, system_prompt: str = "") -> None:
+    def __init__(self, model: str, base_url: str, system_prompt: str = "") -> None:
         super().__init__(system_prompt)
         # API key is retrieved from an environment variable by default
-        self.client = anthropic.Anthropic(max_retries=3, base_url=os.getenv("ANTHROPIC_BASE_URL", "https://api.anthropic.com"))
+        self.client = anthropic.Anthropic(max_retries=3, base_url=base_url)
         self.model = model
 
     def create_messages(self, user_prompt: str) -> List[Dict[str, str]]:
@@ -111,32 +111,31 @@ def get_response(self, response: Dict[str, Any]) -> str:
 
 
 class ChatGPT(LLM):
-    def __init__(self, model: str, system_prompt: str = "") -> None:
+    def __init__(self, model: str, base_url: str, system_prompt: str = "") -> None:
         super().__init__(system_prompt)
-         # Retrieves API key and API Endpoint if specified from an environment variable
-        self.client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_BASE_URL", f"https://api.openai.com/v1"))
+        self.client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=base_url)
         self.model = model
 
     def create_messages(self, user_prompt: str) -> List[Dict[str, str]]:
-        messages = [{"role": "system", "content": self.system_prompt}, {"role": "user", "content": user_prompt}]
+        messages = [{"role": "system", "content": self.system_prompt}, 
+                    {"role": "user", "content": user_prompt}]
         return messages
 
-    def send_message(self, messages: List[Dict[str, str]], max_tokens: int, response_model) -> Dict[str, Any]:
+    def send_message(self, messages: List[Dict[str, str]], max_tokens: int, response_model=None) -> Dict[str, Any]:
         try:
-            # For analyzing files and context code, use the beta endpoint and parse so we can feed it the pydantic model
+            params = {
+                "model": self.model,
+                "messages": messages,
+                "max_tokens": max_tokens,
+            }
+
+            # Add response format configuration if a model is provided
             if response_model:
-                return self.client.beta.chat.completions.parse(
-                    model=self.model,
-                    messages=messages,
-                    max_tokens=max_tokens,
-                    response_format=response_model
-                )
-            else:
-                return self.client.chat.completions.create(
-                    model=self.model,
-                    messages=messages,
-                    max_tokens=max_tokens,
-                )
+                params["response_format"] = {
+                    "type": "json_object"
+                }
+
+            return self.client.chat.completions.create(**params)
         except openai.APIConnectionError as e:
             raise APIConnectionError("The server could not be reached") from e
         except openai.RateLimitError as e:
@@ -146,37 +145,24 @@ def send_message(self, messages: List[Dict[str, str]], max_tokens: int, response
         except Exception as e:
             raise LLMError(f"An unexpected error occurred: {str(e)}") from e
 
-    def _clean_response(self, response: str) -> str:
-        # Step 1: Remove markdown code block wrappers
-        cleaned_text = response.strip('```json\n').strip('```')
-        # Step 2: Correctly handle newlines and escaped characters
-        cleaned_text = cleaned_text.replace('\n', '').replace('\\\'', '\'')
-        # Step 3: Replace escaped double quotes with regular double quotes
-        cleaned_text = cleaned_text.replace('\\"', '"')
-
-        return cleaned_text.replace('\n', '')
-
     def get_response(self, response: Dict[str, Any]) -> str:
         response = response.choices[0].message.content
-        cleaned_response = self._clean_response(response)
-        return cleaned_response
+        return response
 
 
 class Ollama(LLM):
-    def __init__(self, model: str, system_prompt: str = "") -> None:
+    def __init__(self, model: str, base_url: str, system_prompt: str = "") -> None:
         super().__init__(system_prompt)
-        self.api_url = "http://localhost:11434/api/chat"
+        self.api_url = base_url
         self.model = model
 
-    def create_messages(self, user_prompt: str) -> List[Dict[str, str]]:
-        messages = [{"role": "system", "content": self.system_prompt}, {"role": "user", "content": user_prompt}]
-        return messages
+    def create_messages(self, user_prompt: str) -> str:
+        return user_prompt
 
     def send_message(self, user_prompt: str, max_tokens: int, response_model: BaseModel) -> Dict[str, Any]:
         payload = {
             "model": self.model,
-            #"messages": messages,
-            "messages": "hello",
+            "prompt": user_prompt,
             "options": {
             "temperature": 1,
             "system": self.system_prompt,
@@ -195,12 +181,10 @@ def send_message(self, user_prompt: str, max_tokens: int, response_model: BaseMo
             else:
                 raise APIStatusError(e.response.status_code, e.response.json()) from e
 
-
     def get_response(self, response: Dict[str, Any]) -> str:
-        response = response.json()['message']['content']
+        response = response.json()['response']
         return response
 
-
     def _log_response(self, response: Dict[str, Any]) -> None:
         log.debug("Received chat response", extra={"usage": "Ollama"})
 

Diff for: vulnhuntr/__main__.py

@@ -282,17 +282,21 @@ def extract_between_tags(tag: str, string: str, strip: bool = False) -> list[str
     return ext_list
 
 def initialize_llm(llm_arg: str, system_prompt: str = "") -> Claude | ChatGPT | Ollama:
+    llm_arg = llm_arg.lower()
     if llm_arg == 'claude':
         anth_model = os.getenv("ANTHROPIC_MODEL", "claude-3-5-sonnet-20240620")
-        llm = Claude(anth_model, system_prompt)
+        anth_base_url = os.getenv("ANTHROPIC_BASE_URL", "https://api.anthropic.com")
+        llm = Claude(anth_model, anth_base_url, system_prompt)
     elif llm_arg == 'gpt':
         openai_model = os.getenv("OPENAI_MODEL", "gpt-4o-2024-08-06")
-        llm = ChatGPT(openai_model, system_prompt)
+        openai_base_url = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
+        llm = ChatGPT(openai_model, openai_base_url, system_prompt)
     elif llm_arg == 'ollama':
         ollama_model = os.getenv("OLLAMA_MODEL", "llama3")
-        llm = Ollama(ollama_model, system_prompt)
+        ollama_base_url = os.getenv("OLLAMA_BASE_URL", "http://127.0.0.1:11434/api/generate")
+        llm = Ollama(ollama_model, ollama_base_url, system_prompt)
     else:
-        raise ValueError(f"Invalid LLM argument: {llm_arg}")
+        raise ValueError(f"Invalid LLM argument: {llm_arg}\nValid options are: claude, gpt, ollama")
     return llm
 
 def print_readable(report: Response) -> None: