refactor: Improve Ollama integration and add Ollama URL option#9
Merged
refactor: Improve Ollama integration and add Ollama URL option#9
Conversation
This change refactors the Ollama integration to improve performance and adds an option to specify the Ollama URL. The Ollama client is now initialized once and reused, and a new client is created for each process in parallel operations. An --ollama-url argument has been added to the CLI to allow users to specify the URL of their Ollama instance. Additionally, log messages related to Ollama connection have been improved.
Contributor
Reviewer's Guide by SourceryThis pull request refactors the Ollama integration to improve performance by reusing the Ollama client and adds an option to specify the Ollama URL via the CLI. It also improves logging for Ollama connection errors. The Ollama client is now initialized once and reused, and a new client is created for each process in parallel operations. An --ollama-url argument has been added to the CLI to allow users to specify the URL of their Ollama instance. Additionally, log messages related to Ollama connection have been improved. Sequence diagram for getting Ollama clientsequenceDiagram
participant OM as OllamaManager
OM->OM: if not self.client
OM->OM: try
OM->OM: self.client = ollama.Client(self.api_url)
OM->OM: self.client.list()
OM->OM: return self.client
OM->OM: except Exception as e
OM->OM: self._log_connection_error(e)
OM->OM: raise ConnectionError
OM->OM: end
OM->OM: return self.client
Sequence diagram for parallel file processingsequenceDiagram
participant P as Process
participant OM as OllamaManager
participant EF as Embedding Functions
P->OM: ollama_manager = OllamaManager(args.api_url)
P->EF: functions = extract_functions_from_file(args.input_path, content, ollama_manager)
loop For each function
P->EF: func_embedding = generate_content_embedding(func_content, args.embed_model, args.chunk_size, ollama_manager)
end
P->EF: file_embedding = generate_content_embedding(content, args.embed_model, args.chunk_size, ollama_manager)
Updated class diagram for EmbeddingManagerclassDiagram
class EmbeddingManager {
-ollama_manager: OllamaManager
-ollama_client: ollama.Client
__init__(args, ollama_manager: OllamaManager)
index_code_files(files: List[Path]) : None
process_input_files(args) : bool
build_vulnerability_embedding_prompt(vulnerability: Union[str, Dict]) : str
}
Updated class diagram for SecurityAnalyzerclassDiagram
class SecurityAnalyzer {
-llm_model: str
-embedding_manager: EmbeddingManager
-ollama_manager: OllamaManager
-client: ollama.Client
__init__(llm_model: str, embedding_manager: EmbeddingManager, ollama_manager: OllamaManager)
process_analysis_with_model(vulnerabilities, args, report)
analyze_vulnerability(vulnerability: Dict, args, report)
_build_prompt(vulnerability: Dict, code_snippet: str, args) : str
_extract_relevant_info(vulnerability: Dict) : str
}
Updated class diagram for EmbeddingAnalyzerclassDiagram
class EmbeddingAnalyzer {
-embedding_manager: EmbeddingManager
-ollama_manager: OllamaManager
-code_base: dict
-embedding_model: str
__init__(embedding_manager: EmbeddingManager, ollama_manager: OllamaManager)
analyze_all_vulnerabilities(vulnerabilities: List[Dict]) : List[Dict]
analyze_vulnerability(vulnerability: Dict) : Dict
_prepare_analysis_args(vuln: Dict) : list
}
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Contributor
There was a problem hiding this comment.
Hey @psyray - I've reviewed your changes - here's some feedback:
Overall Comments:
- Consider adding a try-except block around the
OllamaManagerinitialization inEmbeddingManagerto handle potential connection errors early. - The
_prepare_analysis_argsfunction inEmbeddingAnalyzerduplicates the Ollama URL, which could lead to inconsistencies; consider passing theOllamaManagerinstance instead.
Here's what I looked at during the review
- 🟡 General issues: 2 issues found
- 🟢 Security: all looks good
- 🟢 Testing: all looks good
- 🟢 Complexity: all looks good
- 🟢 Documentation: all looks good
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
This change enhances the thread safety of the Ollama client initialization and handles cases where the Ollama manager might be missing. It also refactors the _init_ollama method to accept an optional ollama_url argument, improving flexibility. Additionally, explicit checks for a non-null ollama_manager are added before generating content embeddings and extracting functions.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fix #8
This change refactors the Ollama integration to improve performance and adds an option to specify the Ollama URL. The Ollama client is now initialized once and reused, and a new client is created for each process in parallel operations. An --ollama-url argument has been added to the CLI to allow users to specify the URL of their Ollama instance. Additionally, log messages related to Ollama connection have been improved.