Skip to content

style: format code with Black #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
deepsource-autofix wants to merge 1 commit into
base: main
Choose a base branch
from
Open

style: format code with Black #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 16 additions & 10 deletions launch_web.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,37 +7,42 @@
import sys
from pathlib import Path


def main():
"""Launch the VulnForge web interface"""
try:
from streamlit.web import cli as stcli

# Get the UI file path
ui_file = Path(__file__).parent / "modules" / "web" / "dashboard.py"

if not ui_file.exists():
print(f"❌ Error: Web UI file not found at {ui_file}")
print("Please ensure the Robin module is installed correctly.")
return 1

print("🌐 Launching VulnForge Web Interface...")
print("📍 Access the UI at: http://localhost:8501")
print("⚠️ Press Ctrl+C to stop the server\n")

# Prepare streamlit arguments
sys.argv = [
"streamlit",
"run",
str(ui_file),
"--server.port", "8501",
"--server.address", "localhost",
"--server.headless", "true",
"--browser.gatherUsageStats", "false"
"--server.port",
"8501",
"--server.address",
"localhost",
"--server.headless",
"true",
"--browser.gatherUsageStats",
"false",
]

# Launch streamlit
sys.exit(stcli.main())

except ImportError:
print("❌ Error: Streamlit is not installed.")
print("Install it with: pip install streamlit")
Expand All @@ -49,5 +54,6 @@ def main():
print(f"❌ Error launching web interface: {e}")
return 1


if __name__ == "__main__":
sys.exit(main())
134 changes: 69 additions & 65 deletions modules/darkweb/robin/llm_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,21 +14,17 @@

# Fix imports to work both as module and when run directly
try:
from .config import (
OLLAMA_BASE_URL,
OLLAMA_MAIN_MODEL,
OLLAMA_ASSISTANT_MODEL
)
from .config import OLLAMA_BASE_URL, OLLAMA_MAIN_MODEL, OLLAMA_ASSISTANT_MODEL
except ImportError:
from config import (
OLLAMA_BASE_URL,
OLLAMA_MAIN_MODEL,
OLLAMA_ASSISTANT_MODEL
)
from config import OLLAMA_BASE_URL, OLLAMA_MAIN_MODEL, OLLAMA_ASSISTANT_MODEL


class BufferedStreamingHandler(BaseCallbackHandler):
def __init__(self, buffer_limit: int = 60, ui_callback: Optional[Callable[[str], None]] = None):
def __init__(
self,
buffer_limit: int = 60,
ui_callback: Optional[Callable[[str], None]] = None,
):
self.buffer = ""
self.buffer_limit = buffer_limit
self.ui_callback = ui_callback
Expand All @@ -49,7 +45,6 @@ def on_llm_end(self, response, **kwargs) -> None:
self.buffer = ""



def _normalize_model_name(name: str) -> str:
return name.strip().lower()

Expand All @@ -74,67 +69,69 @@ def _get_ollama_base_url() -> str:
# Map input model choices (lowercased) to their configuration
# Each config includes the class and any model-specific constructor parameters
_llm_config_map = {
'gpt-4o': {
'class': ChatOpenAI,
'constructor_params': {'model_name': 'gpt-4o'}
},
'gpt-4.1': {
'class': ChatOpenAI,
'constructor_params': {'model_name': 'gpt-4.1'}
"gpt-4o": {"class": ChatOpenAI, "constructor_params": {"model_name": "gpt-4o"}},
"gpt-4.1": {"class": ChatOpenAI, "constructor_params": {"model_name": "gpt-4.1"}},
"gpt-5.1": {"class": ChatOpenAI, "constructor_params": {"model_name": "gpt-5.1"}},
"gpt-5-mini": {
"class": ChatOpenAI,
"constructor_params": {"model_name": "gpt-5-mini"},
},
'gpt-5.1': {
'class': ChatOpenAI,
'constructor_params': {'model_name': 'gpt-5.1'}
"gpt-5-nano": {
"class": ChatOpenAI,
"constructor_params": {"model_name": "gpt-5-nano"},
},
'gpt-5-mini': {
'class': ChatOpenAI,
'constructor_params': {'model_name': 'gpt-5-mini'}
"claude-3-5-sonnet-latest": {
"class": ChatAnthropic,
"constructor_params": {"model": "claude-3-5-sonnet-latest"},
},
'gpt-5-nano': {
'class': ChatOpenAI,
'constructor_params': {'model_name': 'gpt-5-nano'}
"claude-sonnet-4-5": {
"class": ChatAnthropic,
"constructor_params": {"model": "claude-sonnet-4-5"},
},
'claude-3-5-sonnet-latest': {
'class': ChatAnthropic,
'constructor_params': {'model': 'claude-3-5-sonnet-latest'}
"claude-sonnet-4-0": {
"class": ChatAnthropic,
"constructor_params": {"model": "claude-sonnet-4-0"},
},
'claude-sonnet-4-5': {
'class': ChatAnthropic,
'constructor_params': {'model': 'claude-sonnet-4-5'}
"gemini-2.5-flash": {
"class": ChatGoogleGenerativeAI,
"constructor_params": {"model": "gemini-2.5-flash"},
},
'claude-sonnet-4-0': {
'class': ChatAnthropic,
'constructor_params': {'model': 'claude-sonnet-4-0'}
"gemini-2.5-flash-lite": {
"class": ChatGoogleGenerativeAI,
"constructor_params": {"model": "gemini-2.5-flash-lite"},
},
'gemini-2.5-flash': {
'class': ChatGoogleGenerativeAI,
'constructor_params': {'model': 'gemini-2.5-flash'}
"gemini-2.5-pro": {
"class": ChatGoogleGenerativeAI,
"constructor_params": {"model": "gemini-2.5-pro"},
},
'gemini-2.5-flash-lite': {
'class': ChatGoogleGenerativeAI,
'constructor_params': {'model': 'gemini-2.5-flash-lite'}
"llama3.2": {
"class": ChatOllama,
"constructor_params": {
"model": "llama3.2:latest",
"base_url": _get_ollama_base_url(),
},
},
'gemini-2.5-pro': {
'class': ChatGoogleGenerativeAI,
'constructor_params': {'model': 'gemini-2.5-pro'}
"llama3.1": {
"class": ChatOllama,
"constructor_params": {
"model": "llama3.1:latest",
"base_url": _get_ollama_base_url(),
},
},
'llama3.2': {
'class': ChatOllama,
'constructor_params': {'model': 'llama3.2:latest', 'base_url': _get_ollama_base_url()}
"gemma3": {
"class": ChatOllama,
"constructor_params": {
"model": "gemma3:latest",
"base_url": _get_ollama_base_url(),
},
},
'llama3.1': {
'class': ChatOllama,
'constructor_params': {'model': 'llama3.1:latest', 'base_url': _get_ollama_base_url()}
"deepseek-r1": {
"class": ChatOllama,
"constructor_params": {
"model": "deepseek-r1:latest",
"base_url": _get_ollama_base_url(),
},
},
'gemma3': {
'class': ChatOllama,
'constructor_params': {'model': 'gemma3:latest', 'base_url': _get_ollama_base_url()}
},
'deepseek-r1': {
'class': ChatOllama,
'constructor_params': {'model': 'deepseek-r1:latest', 'base_url': _get_ollama_base_url()}
}

# Add more models here easily:
# 'mistral7b': {
# 'class': ChatOllama,
Expand All @@ -146,6 +143,7 @@ def _get_ollama_base_url() -> str:
# }
}


def fetch_ollama_models() -> List[str]:
"""
Retrieve the list of locally available Ollama models by querying the Ollama HTTP API.
Expand Down Expand Up @@ -178,7 +176,7 @@ def get_model_choices() -> List[str]:
dynamic_models = fetch_ollama_models()

normalized = {_normalize_model_name(m): m for m in base_models}

# Add models from .env if they aren't already there
for env_model in [OLLAMA_MAIN_MODEL, OLLAMA_ASSISTANT_MODEL]:
if env_model:
Expand Down Expand Up @@ -214,15 +212,21 @@ def resolve_model_config(model_choice: str):
if env_model and _normalize_model_name(env_model) == model_choice_lower:
return {
"class": ChatOllama,
"constructor_params": {"model": env_model, "base_url": _get_ollama_base_url()},
"constructor_params": {
"model": env_model,
"base_url": _get_ollama_base_url(),
},
}

# Check against dynamic models from Ollama API
for ollama_model in fetch_ollama_models():
if _normalize_model_name(ollama_model) == model_choice_lower:
return {
"class": ChatOllama,
"constructor_params": {"model": ollama_model, "base_url": _get_ollama_base_url()},
"constructor_params": {
"model": ollama_model,
"base_url": _get_ollama_base_url(),
},
}

return None
return None
26 changes: 22 additions & 4 deletions modules/darkweb/robin/ui.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,12 @@
except ImportError:
# Fall back to absolute imports (when run directly by Streamlit)
import os

# Add parent directory to path
current_dir = Path(__file__).parent
if str(current_dir) not in sys.path:
sys.path.insert(0, str(current_dir))

from scrape import scrape_multiple
from search import get_search_results
from llm_utils import BufferedStreamingHandler, get_model_choices
Expand Down Expand Up @@ -81,7 +82,11 @@ def cached_scrape_multiple(filtered: list, threads: int):
default_model = OLLAMA_MAIN_MODEL or "gpt-5-mini"
default_model_index = (
next(
(idx for idx, name in enumerate(model_options) if name.lower() == default_model.lower()),
(
idx
for idx, name in enumerate(model_options)
if name.lower() == default_model.lower()
),
0,
)
if model_options
Expand All @@ -95,9 +100,20 @@ def cached_scrape_multiple(filtered: list, threads: int):
key="model_select",
)
from .llm_utils import fetch_ollama_models

ollama_reachable = len(fetch_ollama_models()) > 0

if any(name not in {"gpt4o", "gpt-4.1", "claude-3-5-sonnet-latest", "llama3.1", "gemini-2.5-flash"} for name in model_options):
if any(
name
not in {
"gpt4o",
"gpt-4.1",
"claude-3-5-sonnet-latest",
"llama3.1",
"gemini-2.5-flash",
}
for name in model_options
):
if ollama_reachable:
st.sidebar.caption("✅ Locally detected Ollama models are added to this list.")
else:
Expand All @@ -111,7 +127,9 @@ def cached_scrape_multiple(filtered: list, threads: int):
# Main UI - logo and input
_, logo_col, _ = st.columns(3)
with logo_col:
logo_path = Path(__file__).resolve().parent / ".github" / "assets" / "robin_logo.png"
logo_path = (
Path(__file__).resolve().parent / ".github" / "assets" / "robin_logo.png"
)
if logo_path.exists():
st.image(str(logo_path), width=200)

Expand Down
Loading
Loading