feat(llm-proxy): Add support for other claude models + minor fixes

Signed-off-by: Diwank Singh Tomer <[email protected]>
julep-ai · Sep 7, 2024 · 888887a · 888887a
1 parent c54f172
commit 888887a
Show file tree

Hide file tree

Showing 12 changed files with 34 additions and 11 deletions.
diff --git a/agents-api/agents_api/activities/utils.py b/agents-api/agents_api/activities/utils.py
@@ -1,6 +1,8 @@
 import json
 from functools import reduce
 from itertools import accumulate
+from random import random
+from time import time
 from typing import Any, Callable
 
 import re2
@@ -29,13 +31,15 @@
     "match_regex": lambda pattern, string: bool(re2.fullmatch(pattern, string)),
     "max": max,
     "min": min,
+    "random": random,
     "range": range,
     "reduce": reduce,
     "round": round,
     "search_regex": lambda pattern, string: re2.search(pattern, string),
     "set": set,
     "str": str,
     "sum": sum,
+    "time": time,
     "tuple": tuple,
     "zip": zip,
 }

diff --git a/agents-api/agents_api/models/session/get_session.py b/agents-api/agents_api/models/session/get_session.py
@@ -94,7 +94,7 @@ def get_session(
             render_templates,
             token_budget,
             context_overflow,
-            @ "NOW"
+            @ "END"
         }, updated_at = to_int(validity)
     """
 

diff --git a/agents-api/agents_api/models/session/list_sessions.py b/agents-api/agents_api/models/session/list_sessions.py
@@ -101,7 +101,7 @@ def list_sessions(
                 metadata,
                 token_budget,
                 context_overflow,
-                @ "NOW"
+                @ "END"
             }},
             users_p[users, id],
             participants[agents, "agent", id],

diff --git a/agents-api/agents_api/models/session/patch_session.py b/agents-api/agents_api/models/session/patch_session.py
@@ -95,7 +95,7 @@ def patch_session(
             input[{session_update_cols}],
             ids[session_id, developer_id],
             *sessions{{
-                {rest_fields}, metadata: md, @ "NOW"
+                {rest_fields}, metadata: md, @ "END"
             }},
             updated_at = 'ASSERT',
             metadata = concat(md, $metadata),

diff --git a/agents-api/agents_api/models/session/prepare_session_data.py b/agents-api/agents_api/models/session/prepare_session_data.py
@@ -159,7 +159,7 @@ def prepare_session_data(
             render_templates,
             token_budget,
             context_overflow,
-            @ "NOW"
+            @ "END"
         },
         updated_at = to_int(validity),
         record = {

diff --git a/agents-api/agents_api/models/session/update_session.py b/agents-api/agents_api/models/session/update_session.py
@@ -84,7 +84,7 @@ def update_session(
             input[{session_update_cols}],
             ids[session_id, developer_id],
             *sessions{{
-                {rest_fields}, @ "NOW"
+                {rest_fields}, @ "END"
             }},
             updated_at = 'ASSERT'
 

diff --git a/agents-api/agents_api/models/task/get_task.py b/agents-api/agents_api/models/task/get_task.py
@@ -64,7 +64,7 @@ def get_task(
             workflows,
             created_at,
             metadata,
-            @ 'NOW'
+            @ 'END'
         },
         updated_at = to_int(updated_at_ms) / 1000
 

diff --git a/agents-api/agents_api/models/task/list_tasks.py b/agents-api/agents_api/models/task/list_tasks.py
@@ -70,7 +70,7 @@ def list_tasks(
             workflows,
             created_at,
             metadata,
-            @ 'NOW'
+            @ 'END'
         }},
         updated_at = to_int(updated_at_ms) / 1000
 

diff --git a/agents-api/agents_api/models/utils.py b/agents-api/agents_api/models/utils.py
@@ -97,7 +97,7 @@ def mark_session_updated_query(developer_id: UUID | str, session_id: UUID | str)
             render_templates,
             token_budget,
             context_overflow,
-            @ 'NOW'
+            @ 'END'
         }},
         updated_at = [floor(now()), true]
 

diff --git a/llm-proxy/litellm-config.yaml b/llm-proxy/litellm-config.yaml
@@ -38,6 +38,24 @@ model_list:
     tags: ["paid"]
     api_key: os.environ/ANTHROPIC_API_KEY
 
+- model_name: "claude-3-opus"
+  litellm_params:
+    model: "claude-3-opus-20240229"
+    tags: ["paid"]
+    api_key: os.environ/ANTHROPIC_API_KEY
+
+- model_name: "claude-3-sonnet"
+  litellm_params:
+    model: "claude-3-sonnet-20240229"
+    tags: ["paid"]
+    api_key: os.environ/ANTHROPIC_API_KEY
+
+- model_name: "claude-3-haiku"
+  litellm_params:
+    model: "claude-3-haiku-20240307"
+    tags: ["paid"]
+    api_key: os.environ/ANTHROPIC_API_KEY
+
 # Groq models
 - model_name: "llama-3.1-70b"
   litellm_params:

diff --git a/memory-store/Dockerfile b/memory-store/Dockerfile
@@ -39,7 +39,7 @@ RUN \
     rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
 # Set fallback mount directory
-ENV COZO_MNT_DIR=/data COZO_BACKUP_DIR=/backup APP_HOME=/app 
+ENV COZO_MNT_DIR=/data COZO_BACKUP_DIR=/backup APP_HOME=/app COZO_PORT=9070
 WORKDIR $APP_HOME
 
 # Copy the cozo binary

diff --git a/memory-store/backup.sh b/memory-store/backup.sh
@@ -4,11 +4,12 @@ set -eo pipefail  # Exit on error
 set -u  # Exit on undefined variable
 
 # Ensure environment variables are set
-if [ -z "$COZO_PORT" ] || [ -z "$COZO_AUTH_TOKEN" ]; then
-    echo "COZO_PORT or COZO_AUTH_TOKEN is not set"
+if [ -z "$COZO_AUTH_TOKEN" ]; then
+    echo "COZO_AUTH_TOKEN is not set"
     exit 1
 fi
 
+COZO_PORT=${COZO_PORT:-9070}
 COZO_BACKUP_DIR=${COZO_BACKUP_DIR:-/backup}
 TIMESTAMP=$(date +%Y-%m-%d_%H-%M-%S)
 MAX_BACKUPS=${MAX_BACKUPS:-10}