From b5ec5c7996d5230f154a607f66a16f6951a197d1 Mon Sep 17 00:00:00 2001
From: Daniel Chalef <131175+danielchalef@users.noreply.github.com>
Date: Fri, 21 Jun 2024 13:36:50 -0700
Subject: [PATCH] fix/azure openai (#346)

* Fix Azure OpenAI embedding startup issue.

* add new openai models

* cleanup
---
 config.yaml             |  2 ++
 docker-compose.dev.yaml |  1 -
 docker-compose.yaml     | 15 +++++++--------
 pkg/llms/llm_base.go    |  4 ++++
 pkg/testutils/data.go   |  2 --
 pkg/testutils/utils.go  |  2 --
 6 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/config.yaml b/config.yaml
index 71478895..4175486a 100644
--- a/config.yaml
+++ b/config.yaml
@@ -16,6 +16,8 @@ llm:
   # embeddings deployment is required when Zep is configured to use OpenAI embeddings
   #   embedding_deployment: "text-embedding-ada-002-customname"
   # Use only with an alternate OpenAI-compatible API endpoint
+    llm_deployment:
+    embedding_deployment:
   openai_endpoint:
   openai_org_id:
 nlp:
diff --git a/docker-compose.dev.yaml b/docker-compose.dev.yaml
index 077b86d9..d4277ed3 100644
--- a/docker-compose.dev.yaml
+++ b/docker-compose.dev.yaml
@@ -1,4 +1,3 @@
-version: "3.7"
 services:
   db:
     image: ghcr.io/getzep/postgres:latest
diff --git a/docker-compose.yaml b/docker-compose.yaml
index 8ec780a1..c8b9cc32 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -1,10 +1,9 @@
-version: "3.7"
 services:
   db:
     image: ghcr.io/getzep/postgres:latest
     container_name: zep-postgres
     restart: on-failure
-    shm_size: "128mb"  # Increase this if vacuuming fails with a "no space left on device" error
+    shm_size: "128mb" # Increase this if vacuuming fails with a "no space left on device" error
     environment:
       - POSTGRES_USER=postgres
       - POSTGRES_PASSWORD=postgres
@@ -16,7 +15,7 @@ services:
     volumes:
       - zep-db:/var/lib/postgresql/data
     healthcheck:
-      test: [ "CMD", "pg_isready", "-q", "-d", "postgres", "-U", "postgres" ]
+      test: ["CMD", "pg_isready", "-q", "-d", "postgres", "-U", "postgres"]
       interval: 5s
       timeout: 5s
       retries: 5
@@ -39,12 +38,12 @@ services:
     container_name: zep
     restart: on-failure
     depends_on:
-        db:
-          condition: service_healthy
-        nlp:
-          condition: service_healthy
+      db:
+        condition: service_healthy
+      nlp:
+        condition: service_healthy
     ports:
-        - "8000:8000"
+      - "8000:8000"
     volumes:
       - ./config.yaml:/app/config.yaml
     environment:
diff --git a/pkg/llms/llm_base.go b/pkg/llms/llm_base.go
index 117e5145..49a368bd 100644
--- a/pkg/llms/llm_base.go
+++ b/pkg/llms/llm_base.go
@@ -158,6 +158,8 @@ var ValidOpenAILLMs = map[string]bool{
 	"gpt-3.5-turbo-1106": true,
 	"gpt-4-32k":          true,
 	"gpt-4-1106-preview": true,
+	"gpt-4-turbo":        true,
+	"gpt-4o":             true,
 }
 
 var ValidAnthropicLLMs = map[string]bool{
@@ -174,6 +176,8 @@ var MaxLLMTokensMap = map[string]int{
 	"gpt-4":              8192,
 	"gpt-4-32k":          32_768,
 	"gpt-4-1106-preview": 128_000,
+	"gpt-4-turbo":        128_000,
+	"gpt-4o":             128_000,
 	"claude-instant-1":   100_000,
 	"claude-2":           100_000,
 }
diff --git a/pkg/testutils/data.go b/pkg/testutils/data.go
index efb04afc..3586a731 100644
--- a/pkg/testutils/data.go
+++ b/pkg/testutils/data.go
@@ -1,5 +1,3 @@
-//go:build testutils
-
 package testutils
 
 import "github.com/getzep/zep/pkg/models"
diff --git a/pkg/testutils/utils.go b/pkg/testutils/utils.go
index c9a4f3f5..c2f1889b 100644
--- a/pkg/testutils/utils.go
+++ b/pkg/testutils/utils.go
@@ -1,5 +1,3 @@
-//go:build testutils
-
 package testutils
 
 import (