Skip to content

Ollama

OllamaAgent

Bases: BaseAgent

Source code in Docs2KG/agents/ollama.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
class OllamaAgent(BaseAgent):
    def __init__(self, name: str):
        """
        Initialize OllamaAgent with model name and optional API base URL.

        Args:
            name: Name of the Ollama model to use (e.g., 'llama2', 'mistral')
        """
        super().__init__(name)
        self.session = self._init_session()

    def _init_session(self) -> requests.Session:
        """Initialize requests session with retries and timeouts"""
        try:
            session = requests.Session()

            # Configure retries
            retries = Retry(
                total=PROJECT_CONFIG.ollama.max_retries,
                backoff_factor=0.5,
                status_forcelist=[500, 502, 503, 504],
            )

            # Set up the session with retry configuration
            session.mount("http://", HTTPAdapter(max_retries=retries))
            session.mount("https://", HTTPAdapter(max_retries=retries))

            # Set base URL
            self.api_base = PROJECT_CONFIG.ollama.api_base

            logger.info(
                f"Successfully initialized Ollama session for model {self.name}"
            )
            return session

        except Exception as e:
            logger.error(f"Failed to initialize Ollama session: {str(e)}")
            raise

    def reset_session(self):
        """Reset the requests session"""
        self.session.close()
        self.session = self._init_session()

    def process(self, input_data: Any, reset_session: bool = False) -> Any:
        """
        Process input using the Ollama API.

        Args:
            input_data: The input to be processed by the model
            reset_session: Whether to reset the requests session before making the API call

        Returns:
            Dict containing the model response and metadata
        """
        logger.info(f"Using Ollama model: {self.name}")

        try:
            # Prepare the request
            url = f"{self.api_base}/api/generate"

            payload = {
                "model": self.name,
                "prompt": str(input_data),
                "temperature": PROJECT_CONFIG.ollama.temperature,
                "stream": False,
                "format": PROJECT_CONFIG.ollama.format,
            }

            # Make the API call
            if reset_session:
                self.reset_session()
            response = self.session.post(
                url, json=payload, timeout=PROJECT_CONFIG.ollama.timeout
            )
            response.raise_for_status()

            result = response.json()

            return {
                "model": self.name,
                "input": input_data,
                "status": "processed",
                "response": result.get("response", ""),
                "usage": {
                    "eval_count": result.get("eval_count", 0),
                    "eval_duration": result.get("eval_duration", 0),
                    "total_duration": result.get("total_duration", 0),
                },
            }

        except requests.exceptions.RequestException as e:
            logger.error(f"Error making request to Ollama API: {str(e)}")
            raise
        except Exception as e:
            logger.error(f"Error processing input with Ollama: {str(e)}")
            raise

__init__(name)

Initialize OllamaAgent with model name and optional API base URL.

Parameters:

Name Type Description Default
name str

Name of the Ollama model to use (e.g., 'llama2', 'mistral')

 required
Source code in Docs2KG/agents/ollama.py 
12
13
14
15
16
17
18
19
20
def __init__(self, name: str):
    """
    Initialize OllamaAgent with model name and optional API base URL.

    Args:
        name: Name of the Ollama model to use (e.g., 'llama2', 'mistral')
    """
    super().__init__(name)
    self.session = self._init_session()

process(input_data, reset_session=False)

Process input using the Ollama API.

Parameters:

Name Type Description Default
input_data Any

The input to be processed by the model

 required
reset_session bool

Whether to reset the requests session before making the API call

 False

Returns:

Type Description
Any

Dict containing the model response and metadata
Source code in Docs2KG/agents/ollama.py 
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def process(self, input_data: Any, reset_session: bool = False) -> Any:
    """
    Process input using the Ollama API.

    Args:
        input_data: The input to be processed by the model
        reset_session: Whether to reset the requests session before making the API call

    Returns:
        Dict containing the model response and metadata
    """
    logger.info(f"Using Ollama model: {self.name}")

    try:
        # Prepare the request
        url = f"{self.api_base}/api/generate"

        payload = {
            "model": self.name,
            "prompt": str(input_data),
            "temperature": PROJECT_CONFIG.ollama.temperature,
            "stream": False,
            "format": PROJECT_CONFIG.ollama.format,
        }

        # Make the API call
        if reset_session:
            self.reset_session()
        response = self.session.post(
            url, json=payload, timeout=PROJECT_CONFIG.ollama.timeout
        )
        response.raise_for_status()

        result = response.json()

        return {
            "model": self.name,
            "input": input_data,
            "status": "processed",
            "response": result.get("response", ""),
            "usage": {
                "eval_count": result.get("eval_count", 0),
                "eval_duration": result.get("eval_duration", 0),
                "total_duration": result.get("total_duration", 0),
            },
        }

    except requests.exceptions.RequestException as e:
        logger.error(f"Error making request to Ollama API: {str(e)}")
        raise
    except Exception as e:
        logger.error(f"Error processing input with Ollama: {str(e)}")
        raise

reset_session()

Reset the requests session

Source code in Docs2KG/agents/ollama.py 
50
51
52
53
def reset_session(self):
    """Reset the requests session"""
    self.session.close()
    self.session = self._init_session()