Skip to content

Manager

AgentManager

Source code in Docs2KG/agents/manager.py
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
class AgentManager:
    def __init__(self, agent_name: str, agent_type: str, **kwargs):
        """
        Initialize AgentManager with a specific agent.

        Args:
            agent_name: Name for the agent (e.g., 'gpt-4', 'gpt-4-turbo')
            agent_type: Type of agent ('cloud', 'quantization', or 'hf')
        """
        self.agent_types = {
            "cloud": CloudAgent,
            "quantization": QuantizationAgent,
            "ollama": OllamaAgent,
            "hf": HuggingFaceAgent,
        }
        self.agent_type = agent_type

        self.agent = self._init_agent(agent_name, agent_type, **kwargs)

    def _init_agent(self, agent_name: str, agent_type: str, **kwargs) -> BaseAgent:
        agent_type = agent_type.lower()
        if agent_type not in self.agent_types:
            raise InvalidAgentType(
                f"Invalid agent type. Must be one of: {', '.join(self.agent_types.keys())}"
            )

        agent_class = self.agent_types[agent_type]
        return agent_class(agent_name, **kwargs)

    def process_input(self, input_data: Any, reset_session: bool = False) -> Any:
        if self.agent_type == "ollama":
            return self.agent.process(input_data, reset_session)
        return self.agent.process(input_data)

    def get_agent_info(self) -> Dict[str, str]:
        return {
            "name": self.agent.name,
            "type": type(self.agent).__name__,
            "config": getattr(self.agent, "model", None),
        }

__init__(agent_name, agent_type, **kwargs)

Initialize AgentManager with a specific agent.

Parameters:

Name Type Description Default
agent_name str

Name for the agent (e.g., 'gpt-4', 'gpt-4-turbo')

required
agent_type str

Type of agent ('cloud', 'quantization', or 'hf')

required
Source code in Docs2KG/agents/manager.py
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def __init__(self, agent_name: str, agent_type: str, **kwargs):
    """
    Initialize AgentManager with a specific agent.

    Args:
        agent_name: Name for the agent (e.g., 'gpt-4', 'gpt-4-turbo')
        agent_type: Type of agent ('cloud', 'quantization', or 'hf')
    """
    self.agent_types = {
        "cloud": CloudAgent,
        "quantization": QuantizationAgent,
        "ollama": OllamaAgent,
        "hf": HuggingFaceAgent,
    }
    self.agent_type = agent_type

    self.agent = self._init_agent(agent_name, agent_type, **kwargs)