Bring Your Own Model

1
import os
2
from smallestai.atoms.agent.nodes import OutputAgentNode
3
from smallestai.atoms.agent.clients.openai import OpenAIClient
4
from smallestai.atoms.agent.server import AtomsApp
5
from smallestai.atoms.agent.session import AgentSession
6
7
class LocalAgent(OutputAgentNode):
8
    def __init__(self):
9
        super().__init__(name="local-agent")
10
        
11
        # Connect to your local model
12
        self.llm = OpenAIClient(
13
            model="llama3",
14
            base_url="http://localhost:11434/v1",
15
            api_key="ollama"  # Not required for Ollama
16
        )
17
        
18
        self.context.add_message(
19
            "system",
20
            "You are a helpful assistant running on local hardware."
21
        )
22
23
    async def generate_response(self):
24
        response = await self.llm.chat(
25
            messages=self.context.messages,
26
            stream=True
27
        )
28
        async for chunk in response:
29
            if chunk.content:
30
                yield chunk.content
31
32
async def on_start(session: AgentSession):
33
    session.add_node(LocalAgent())
34
    await session.start()
35
    await session.wait_until_complete()
36
37
if __name__ == "__main__":
38
    app = AtomsApp(setup_handler=on_start)
39
    app.run()

1
from smallestai.atoms.agent.clients.openai import OpenAIClient
2
3
llm = OpenAIClient(
4
    model="your-model-name",
5
    base_url="https://your-server.example.com/v1",
6
    api_key=os.getenv("YOUR_API_KEY")
7
)

1
from smallestai.atoms.agent.clients.openai import OpenAIClient
2
3
llm = OpenAIClient(
4
    model="llama3",
5
    base_url="http://localhost:11434/v1",
6
    api_key="ollama"  # Doesn't require a real key
7
)

1
from smallestai.atoms.agent.clients.openai import OpenAIClient
2
3
llm = OpenAIClient(
4
    model="meta-llama/Llama-3-8B-Instruct",
5
    base_url="http://localhost:8000/v1",
6
    api_key="vllm"
7
)

1
from smallestai.atoms.agent.clients.openai import OpenAIClient
2
3
llm = OpenAIClient(
4
    model="local-model",
5
    base_url="http://localhost:1234/v1",
6
    api_key="lmstudio"
7
)