Spaces:
Runtime error
Runtime error
File size: 2,075 Bytes
c6421b9 44c6641 c6421b9 9cb712d c6421b9 44c6641 c6421b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | import types
from langchain_groq import ChatGroq
from langmem import create_memory_store_manager
from app.schemas.memory_agent_schema import EmailMemory
from app.persistance.memory_store_checkpointer_config import memory_store
import os
from app.core.config import settings
if "GROQ_API_KEY" not in os.environ:
os.environ["GROQ_API_KEY"] = settings.GROQ_API_KEY
def patch_groq_for_extractions(model: ChatGroq):
# Capture the original method
original_bind_tools = model.bind_tools
def fixed_bind_tools(self, tools, **kwargs):
fixed_tools = []
for tool in tools:
# 1. Handle dictionary-style tools (JSON Schema)
if isinstance(tool, dict):
# Ensure nested description is present
if "function" in tool and not tool["function"].get("description"):
tool["function"]["description"] = "Extract information based on the schema."
elif not tool.get("description"):
tool["description"] = "Extract information based on the schema."
fixed_tools.append(tool)
# 2. Handle Class-style tools (Pydantic)
else:
if not getattr(tool, "__doc__", None):
# This fixes the internal 'Done' class from langmem/trustcall
tool.__doc__ = "Signal that extraction is finished."
fixed_tools.append(tool)
return original_bind_tools(fixed_tools, **kwargs)
object.__setattr__(
model,
"bind_tools",
types.MethodType(fixed_bind_tools, model)
)
return model
# Apply the patch
model=ChatGroq(model="openai/gpt-oss-20b", temperature=0.2)
model = patch_groq_for_extractions(model)
namespace = ("emails", "{user_id}", "collection")
memory_manager_agent = create_memory_store_manager(
model,
schemas=[EmailMemory],
namespace=namespace,
store=memory_store,
instructions="Extract required info from incoming mail and its reply .",
enable_inserts=True,
enable_deletes=True,
) |