fix: Apply various small fixes

julep-ai · May 3, 2024 · d762fa1 · d762fa1
1 parent 68ce7f5
commit d762fa1
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/agents-api/agents_api/clients/embed.py b/agents-api/agents_api/clients/embed.py
@@ -5,6 +5,7 @@
 async def embed(
     inputs: list[str],
     join_inputs=False,
+    embed_model_name: str = embedding_model_id,
 ) -> list[list[float]]:
     async with httpx.AsyncClient(timeout=30) as client:
         resp = await client.post(
@@ -17,7 +18,7 @@ async def embed(
                 "normalize": True,
                 # FIXME: We should control the truncation ourselves and truncate before sending
                 "truncate": truncate_embed_text,
-                "model_id": embedding_model_id,
+                "model_id": embed_model_name,
             },
         )
         resp.raise_for_status()

diff --git a/agents-api/agents_api/embed_models_registry.py b/agents-api/agents_api/embed_models_registry.py
@@ -78,7 +78,7 @@ async def embed(
         embeddings: list[np.ndarray | list[float]] = []
 
         if self.embedding_provider == "julep":
-            embeddings = await embed(input)
+            embeddings = await embed(input, embed_model_name=self.embedding_model_name)
         elif self.embedding_provider == "openai":
             embeddings = (
                 await openai_client.embeddings.create(

diff --git a/agents-api/agents_api/routers/agents/routers.py b/agents-api/agents_api/routers/agents/routers.py
@@ -3,6 +3,7 @@
 from typing import Annotated
 from uuid import uuid4
 
+from agents_api.autogen.openapi_model import ContentItem
 from agents_api.model_registry import validate_configuration
 from fastapi import APIRouter, HTTPException, status, Depends
 import pandas as pd
@@ -304,7 +305,13 @@ async def list_agents(
 @router.post("/agents/{agent_id}/docs", tags=["agents"])
 async def create_docs(agent_id: UUID4, request: CreateDoc) -> ResourceCreatedResponse:
     doc_id = uuid4()
-    content = [request.content] if isinstance(request.content, str) else request.content
+    content = [
+        (c.model_dump() if isinstance(c, ContentItem) else c)
+        for c in (
+            [request.content] if isinstance(request.content, str) else request.content
+        )
+    ]
+
     resp: pd.DataFrame = create_docs_query(
         owner_type="agent",
         owner_id=agent_id,