Refactoring: now using the langchain class

core/cat/factory/custom_llm.py

@@ -3,7 +3,6 @@
 import requests
 from langchain.llms.base import LLM
 from langchain.llms.openai import OpenAI
-from langchain.llms.ollama import Ollama
 
 
 class LLMDefault(LLM):
@@ -88,18 +87,5 @@ def __init__(self, **kwargs):
         self.url = kwargs['url']
         self.openai_api_base = os.path.join(self.url, "v1")
 
-class CustomOllama(Ollama):
 
-
-    def __init__(self, **kwargs):
-
-        super().__init__(
-            base_url=kwargs["base_url"],
-            model=kwargs["model"],
-            num_ctx=kwargs["num_ctx"],
-            repeat_last_n=kwargs["repeat_last_n"],
-            repeat_penalty=kwargs["repeat_penalty"],
-            temperature=kwargs["temperature"],
-        )
-
 

core/cat/factory/llm.py

@@ -1,12 +1,12 @@
 import langchain
 from langchain.chat_models import ChatOpenAI, AzureChatOpenAI
 from langchain.llms import OpenAI, AzureOpenAI
-
+from langchain.llms.ollama import Ollama
 from typing import Dict, List, Type
 import json
 from pydantic import BaseModel, ConfigDict
 
-from cat.factory.custom_llm import CustomOllama, LLMDefault, LLMCustom, CustomOpenAI
+from cat.factory.custom_llm import LLMDefault, LLMCustom, CustomOpenAI
 
 
 # Base class to manage LLM configuration.
@@ -279,8 +279,8 @@ class LLMCustomOllama(LLMSettings):
     repeat_last_n: int = 64
     repeat_penalty: float = 1.1
     temperature: float = 0.8
-
-    _pyclass: Type = CustomOllama
+    
+    _pyclass: Type = Ollama
 
     model_config = ConfigDict(
         json_schema_extra = {