Add support for Ollama

app/core/config_provider.py

@@ -13,13 +13,21 @@ def __init__(self):
             "password": os.getenv("NEO4J_PASSWORD"),
         }
         self.github_key = os.getenv("GITHUB_PRIVATE_KEY")
+        self.ollama_endpoint = os.getenv("OLLAMA_ENDPOINT", "http://localhost:11434")
+        self.ollama_model = os.getenv("OLLAMA_MODEL", "llama2")
 
     def get_neo4j_config(self):
         return self.neo4j_config
 
     def get_github_key(self):
         return self.github_key
 
+    def get_ollama_config(self):
+        return {
+            "endpoint": self.ollama_endpoint,
+            "model": self.ollama_model,
+        }
+
     def get_demo_repo_list(self):
         return [
             {

app/modules/intelligence/agents/agent_factory.py

@@ -24,6 +24,7 @@
     AgentType,
     ProviderService,
 )
+from langchain_ollama import Ollama
 
 
 class AgentFactory:
@@ -70,6 +71,10 @@ def _create_agent(
             "code_generation_agent": lambda: CodeGenerationChatAgent(
                 mini_llm, reasoning_llm, self.db
             ),
+            "ollama_agent": lambda: Ollama(
+                base_url=self.provider_service.get_ollama_endpoint(),
+                model=self.provider_service.get_ollama_model(),
+            ),
         }
 
         if agent_id in agent_map:

app/modules/intelligence/agents/agent_injector_service.py

@@ -28,6 +28,7 @@
     AgentType,
     ProviderService,
 )
+from langchain_ollama import Ollama
 
 logger = logging.getLogger(__name__)
 
@@ -59,6 +60,10 @@ def _initialize_agents(self) -> Dict[str, Any]:
             "code_generation_agent": CodeGenerationChatAgent(
                 mini_llm, reasoning_llm, self.sql_db
             ),
+            "ollama_agent": Ollama(
+                base_url=self.provider_service.get_ollama_endpoint(),
+                model=self.provider_service.get_ollama_model(),
+            ),
         }
 
     def get_agent(self, agent_id: str) -> Any:

app/modules/intelligence/provider/provider_service.py

@@ -7,6 +7,7 @@
 from langchain_anthropic import ChatAnthropic
 from langchain_openai.chat_models import ChatOpenAI
 from portkey_ai import PORTKEY_GATEWAY_URL, createHeaders
+from langchain_ollama import Ollama
 
 from app.modules.key_management.secret_manager import SecretManager
 from app.modules.users.user_preferences_model import UserPreferences
@@ -44,6 +45,11 @@ async def list_available_llms(self) -> List[ProviderInfo]:
                 name="Anthropic",
                 description="An AI safety-focused company known for models like Claude.",
             ),
+            ProviderInfo(
+                id="ollama",
+                name="Ollama",
+                description="A provider for running open source models locally.",
+            ),
         ]
 
     async def set_global_ai_provider(self, user_id: str, provider: str):
@@ -195,6 +201,12 @@ def get_large_llm(self, agent_type: AgentType):
                         default_headers=portkey_headers,
                     )
 
+        elif preferred_provider == "ollama":
+            logging.info("Initializing Ollama LLM")
+            ollama_endpoint = os.getenv("OLLAMA_ENDPOINT", "http://localhost:11434")
+            ollama_model = os.getenv("OLLAMA_MODEL", "llama2")
+            self.llm = Ollama(base_url=ollama_endpoint, model=ollama_model)
+
         else:
             raise ValueError("Invalid LLM provider selected.")
 
@@ -323,6 +335,12 @@ def get_small_llm(self, agent_type: AgentType):
                         default_headers=portkey_headers,
                     )
 
+        elif preferred_provider == "ollama":
+            logging.info("Initializing Ollama LLM")
+            ollama_endpoint = os.getenv("OLLAMA_ENDPOINT", "http://localhost:11434")
+            ollama_model = os.getenv("OLLAMA_MODEL", "llama2")
+            self.llm = Ollama(base_url=ollama_endpoint, model=ollama_model)
+
         else:
             raise ValueError("Invalid LLM provider selected.")
 
@@ -337,6 +355,8 @@ def get_llm_provider_name(self) -> str:
             return "OpenAI"
         elif isinstance(llm, ChatAnthropic):
             return "Anthropic"
+        elif isinstance(llm, Ollama):
+            return "Ollama"
         elif isinstance(llm, LLM):
             return "OpenAI" if llm.model.split("/")[0] == "openai" else "Anthropic"
         else:

app/modules/intelligence/tools/tool_service.py

@@ -36,6 +36,8 @@
     GetNodesFromTags,
 )
 from app.modules.intelligence.tools.tool_schema import ToolInfo
+from langchain_ollama import Ollama
+from app.core.config_provider import config_provider
 
 
 class ToolService:
@@ -65,8 +67,18 @@ def _initialize_tools(self) -> Dict[str, Any]:
             "get_node_neighbours_from_node_id": GetNodeNeighboursFromNodeIdTool(
                 self.db
             ),
+            "ollama_tool": Ollama(
+                base_url=self._get_ollama_endpoint(),
+                model=self._get_ollama_model(),
+            ),
         }
 
+    def _get_ollama_endpoint(self) -> str:
+        return config_provider.get_ollama_config()["endpoint"]
+
+    def _get_ollama_model(self) -> str:
+        return config_provider.get_ollama_config()["model"]
+
     async def run_tool(self, tool_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
         tool = self.tools.get(tool_id)
         if not tool: