google gemini support, docker startup fix, console input fix
frdel committed
Jul 8, 2024 at 13:59 UTC
76ff3c5fcf1e28bc2ffb9555a6a79ba74f880bef
5 files changed
+48
-13
main.py
+28
-6
@@ -1,10 +1,13 @@
1
+import signal
2
+import sys
3
+import termios
4
import threading, time, models, os
5
+import tty
6
from ansio import application_keypad, mouse_input, raw_input
7
from ansio.input import InputEvent, get_input_event
8
from agent import Agent, AgentConfig
9
from python.helpers.print_style import PrintStyle
10
from python.helpers.files import read_file
7
-from pytimedinput import timedInput as timed_input
11
from python.helpers import files
12
13
@@ -17,7 +20,7 @@ def initialize():
20
# chat_llm = models.get_groq_llama70b(temperature=0.2)
21
# chat_llm = models.get_groq_llama70b_json(temperature=0.2)
22
# chat_llm = models.get_groq_llama8b(temperature=0.2)
20
- chat_llm = models.get_openai_gpt35(temperature=0)
23
+ # chat_llm = models.get_openai_gpt35(temperature=0)
24
# chat_llm = models.get_openai_gpt4o(temperature=0)
25
# chat_llm = models.get_anthropic_opus(temperature=0)
26
# chat_llm = models.get_anthropic_sonnet(temperature=0)
@@ -28,6 +31,7 @@ def initialize():
31
# chat_llm = models.get_ollama(model_name="llama3:8b-text-fp16")
32
# chat_llm = models.get_ollama(model_name="gemma2:latest")
33
# chat_llm = models.get_ollama(model_name="qwen:14b")
34
+ chat_llm = models.get_google_chat()
35
36
# embedding model used for memory
37
# embedding_llm = models.get_embedding_openai()
@@ -78,18 +82,21 @@ def chat(agent:Agent):
82
timeout = agent.get_data("timeout") # how long the agent is willing to wait
83
if not timeout: # if agent wants to wait for user input forever
84
PrintStyle(background_color="#6C3483", font_color="white", bold=True, padding=True).print(f"User message ('exit' to leave):")
85
+ import readline # this fixes arrow keys in terminal
86
user_input = input("> ")
87
PrintStyle(font_color="white", padding=False, log_only=True).print(f"> {user_input}")
88
89
else: # otherwise wait for user input with a timeout
90
PrintStyle(background_color="#6C3483", font_color="white", bold=True, padding=True).print(f"User message ({timeout}s timeout, 'wait' to wait, 'exit' to leave):")
86
- user_input = timed_input("> ", timeout=timeout)
91
+ import readline # this fixes arrow keys in terminal
92
+ # user_input = timed_input("> ", timeout=timeout)
93
+ user_input = timeout_input("> ", timeout=timeout)
94
88
- if user_input[1]:
95
+ if not user_input:
96
user_input = read_file("prompts/fw.msg_timeout.md")
97
PrintStyle(font_color="white", padding=False).stream(f"{user_input}")
98
else:
92
- user_input = user_input[0].strip()
99
+ user_input = user_input.strip()
100
if user_input.lower()=="wait": # the user needs more time
101
user_input = input("> ").strip()
102
PrintStyle(font_color="white", padding=False, log_only=True).print(f"> {user_input}")
@@ -113,7 +120,7 @@ def intervention():
120
Agent.paused = True # stop agent streaming
121
PrintStyle(background_color="#6C3483", font_color="white", bold=True, padding=True).print(f"User intervention ('exit' to leave, empty to continue):")
122
116
- import readline
123
+ import readline # this fixes arrow keys in terminal
124
user_input = input("> ").strip()
125
PrintStyle(font_color="white", padding=False, log_only=True).print(f"> {user_input}")
126
@@ -139,6 +146,21 @@ def capture_keys():
146
intervent=True
147
continue
148
149
+# User input with timeout
150
+def timeout_input(prompt, timeout=10):
151
+ def alarm_handler(signum, frame):
152
+ raise TimeoutError()
153
+
154
+ signal.signal(signal.SIGALRM, alarm_handler)
155
+ signal.alarm(timeout)
156
+
157
+ try:
158
+ return input(prompt)
159
+ except TimeoutError:
160
+ return ""
161
+ finally:
162
+ signal.alarm(0) # Cancel the alarm
163
+
164
if __name__ == "__main__":
165
print("Initializing framework...")
166
models.py
+5
@@ -5,6 +5,7 @@ from langchain_openai import ChatOpenAI, OpenAI, OpenAIEmbeddings
5
from langchain_anthropic import ChatAnthropic
6
from langchain_groq import ChatGroq
7
from langchain_huggingface import HuggingFaceEmbeddings
8
+from langchain_google_genai import ChatGoogleGenerativeAI
9
10
11
# Load environment variables
@@ -81,6 +82,10 @@ def get_ollama_dolphin(api_key=None, temperature=DEFAULT_TEMPERATURE):
82
def get_ollama_phi(api_key=None, temperature=DEFAULT_TEMPERATURE):
83
return Ollama(model="phi3:3.8b-mini-instruct-4k-fp16",temperature=temperature)
84
85
+def get_google_chat(model_name="gemini-1.5-flash-latest", api_key=None, temperature=DEFAULT_TEMPERATURE):
86
+ api_key = api_key or get_api_key("google")
87
+ return ChatGoogleGenerativeAI(model=model_name, temperature=temperature, google_api_key=api_key) # type: ignore
88
+
89
def get_embedding_hf(model_name="sentence-transformers/all-MiniLM-L6-v2"):
90
return HuggingFaceEmbeddings(model_name=model_name)
91
python/helpers/docker.py
+9
-7
@@ -5,7 +5,7 @@ from typing import Dict, Optional
5
from python.helpers.files import get_abs_path
6
7
class DockerContainerManager:
8
- def __init__(self, image:str, name:str, ports: Optional[Dict[str, int]] = None, volumes: Optional[Dict[str, Dict[str, str]]] = None):
8
+ def __init__(self, image: str, name: str, ports: Optional[Dict[str, int]] = None, volumes: Optional[Dict[str, Dict[str, str]]] = None):
9
self.client = docker.from_env()
10
self.image = image
11
self.name = name
@@ -13,8 +13,6 @@ class DockerContainerManager:
13
self.volumes = volumes
14
self.container = None
15
16
-
17
-
16
def cleanup_container(self) -> None:
17
if self.container:
18
try:
@@ -26,14 +24,19 @@ class DockerContainerManager:
24
25
def start_container(self) -> None:
26
existing_container = None
29
- for container in self.client.containers.list():
27
+ for container in self.client.containers.list(all=True):
28
if container.name == self.name:
29
existing_container = container
30
break
31
32
if existing_container:
35
- #print(f"Container with name '{self.name}' is already running with ID: {existing_container.id}")
36
- pass
33
+ if existing_container.status != 'running':
34
+ print(f"Starting existing container: {self.name} for safe code execution...")
35
+ existing_container.start()
36
+ self.container = existing_container
37
+ else:
38
+ self.container = existing_container
39
+ # print(f"Container with name '{self.name}' is already running with ID: {existing_container.id}")
40
else:
41
print(f"Initializing docker container {self.name} for safe code execution...")
42
self.container = self.client.containers.run(
@@ -46,4 +49,3 @@ class DockerContainerManager:
49
atexit.register(self.cleanup_container)
50
print(f"Started container with ID: {self.container.id}")
51
time.sleep(1) # this helps to get SSH ready
49
-
requirements.txt
+1
@@ -6,6 +6,7 @@ langchain-openai==0.1.8
6
langchain-community==0.2.4
7
langchain-anthropic==0.1.15
8
langchain-chroma==0.1.1
9
+google-generativeai==0.5.4
10
webcolors==24.6.0
11
sentence-transformers==3.0.1
12
pytimedinput==2.0.1
test.py
new
+5
@@ -0,0 +1,5 @@
1
+from models import get_google_chat
2
+
3
+llm = get_google_chat()
4
+result = llm.invoke("Write a ballad about LangChain")
5
+print(result.content)