Merge pull request #185 from mindsdb/fix/add-v2-completion_stream_v2

fix: add v2 for methods to return raw results

Author: lucas-koontz

mindsdb_sdk/__about__.py

@@ -1,6 +1,6 @@
 __title__ = 'mindsdb_sdk'
 __package_name__ = 'mindsdb_sdk'
-__version__ = '3.4.0'
+__version__ = '3.4.1'
 __description__ = "MindsDB Python SDK, provides an SDK to use a remote mindsdb instance"
 __email__ = "jorge@mindsdb.com"
 __author__ = 'MindsDB Inc'

mindsdb_sdk/agents.py

@@ -35,7 +35,7 @@ def __repr__(self):
 
 class Agent:
     """Represents a MindsDB agent.
-    
+
     Working with agents:
 
     Get an agent by name:
@@ -100,9 +100,15 @@ def __init__(
     def completion(self, messages: List[dict]) -> AgentCompletion:
         return self.collection.completion(self.name, messages)
 
+    def completion_v2(self, messages: List[dict]) -> AgentCompletion:
+        return self.collection.completion_v2(self.name, messages)
+
     def completion_stream(self, messages: List[dict]) -> Iterable[object]:
         return self.collection.completion_stream(self.name, messages)
 
+    def completion_stream_v2(self, messages: List[dict]) -> Iterable[object]:
+        return self.collection.completion_stream_v2(self.name, messages)
+
     def add_files(self, file_paths: List[str], description: str, knowledge_base: str = None):
         """
         Add a list of files to the agent for retrieval.
@@ -246,6 +252,17 @@ def completion(self, name: str, messages: List[dict]) -> AgentCompletion:
 
         return AgentCompletion(data['message']['content'])
 
+    def completion_v2(self, name: str, messages: List[dict]) -> AgentCompletion:
+        """
+        Queries the agent for a completion.
+
+        :param name: Name of the agent
+        :param messages: List of messages to be sent to the agent
+
+        :return: completion from querying the agent
+        """
+        return self.api.agent_completion(self.project.name, name, messages)
+
     def completion_stream(self, name, messages: List[dict]) -> Iterable[object]:
         """
         Queries the agent for a completion and streams the response as an iterable object.
@@ -257,6 +274,17 @@ def completion_stream(self, name, messages: List[dict]) -> Iterable[object]:
         """
         return self.api.agent_completion_stream(self.project.name, name, messages)
 
+    def completion_stream_v2(self, name, messages: List[dict]) -> Iterable[object]:
+        """
+        Queries the agent for a completion and streams the response as an iterable object.
+
+        :param name: Name of the agent
+        :param messages: List of messages to be sent to the agent
+
+        :return: iterable of completion chunks from querying the agent.
+        """
+        return self.api.agent_completion_stream_v2(self.project.name, name, messages)
+
     def _create_default_knowledge_base(self, agent: Agent, name: str) -> KnowledgeBase:
         # Make sure default ML engine for embeddings exists.
         try:

mindsdb_sdk/connectors/rest_api.py

@@ -37,7 +37,7 @@ def _raise_for_status(response):
 
 
 class RestAPI:
-    def __init__(self, url=None, login=None, password=None, api_key=None, is_managed=False, 
+    def __init__(self, url=None, login=None, password=None, api_key=None, is_managed=False,
                  cookies=None, headers=None):
 
         self.url = url
@@ -289,6 +289,22 @@ def agent_completion_stream(self, project: str, name: str, messages: List[dict])
             # Stream objects loaded from SSE events 'data' param.
             yield json.loads(chunk.data)
 
+    @_try_relogin
+    def agent_completion_stream_v2(self, project: str, name: str, messages: List[dict]):
+        url = self.url + f'/api/projects/{project}/agents/{name}/completions/stream'
+        response = self.session.post(url, json={'messages': messages}, stream=True)
+
+        # Check for HTTP errors before processing the stream
+        response.raise_for_status()
+
+        client = SSEClient(response)
+
+        try:
+            for chunk in client.events():
+                yield chunk  # Stream SSE events
+        except Exception as e:
+            yield e
+
     @_try_relogin
     def create_agent(self, project: str, name: str, model: str = None, provider: str = None, skills: List[str] = None, params: dict = None):
         url = self.url + f'/api/projects/{project}/agents'
@@ -452,4 +468,3 @@ def knowledge_base_completion(self, project: str, knowledge_base_name, payload):
         )
         _raise_for_status(r)
         return r.json()
-