Merge pull request #141 from mindsdb/staging

Release 3.0.2

examples/using_agents_with_retrieval.py

@@ -9,11 +9,14 @@
 
 # Now create an agent that will use the model we just created.
 agent = con.agents.create(name=f'mindsdb_retrieval_agent_{model_name}_{uuid4().hex}',
-                          model='gpt-4')
+                          model='gpt-4',
+                          params={'return_context': True})
 
 agent.add_file('./data/tokaido-rulebook.pdf', 'rule book for the board game Tokaido')
 
 
 question = "what are the rules for the game takaido?"
 answer = agent.completion([{'question': question, 'answer': None}])
-print(answer.content)
+print(answer.context)
+print(answer)
+

mindsdb_sdk/__about__.py

@@ -1,6 +1,6 @@
 __title__ = 'mindsdb_sdk'
 __package_name__ = 'mindsdb_sdk'
-__version__ = '3.0.1'
+__version__ = '3.0.2'
 __description__ = "MindsDB Python SDK, provides an SDK to use a remote mindsdb instance"
 __email__ = "[email protected]"
 __author__ = 'MindsDB Inc'

mindsdb_sdk/agents.py

@@ -16,12 +16,21 @@
 _DEFAULT_LLM_MODEL = 'gpt-4o'
 
 class AgentCompletion:
-    """Represents a full MindsDB agent completion"""
-    def __init__(self, content: str):
+    """
+    Represents a full MindsDB agent completion response.
+
+    Attributes:
+    content: The completion content.
+    context: Only relevant for retrieval agents. Contains the context retrieved from the knowledge base.
+
+
+    """
+    def __init__(self, content: str, context: List[dict] = None):
         self.content = content
+        self.context = context
 
     def __repr__(self):
-        return self.content
+        return f'{self.__class__.__name__}(content: {self.content}, context: {self.context})'
 
 
 class Agent:
@@ -208,6 +217,9 @@ def completion(self, name: str, messages: List[dict]) -> AgentCompletion:
         :return: completion from querying the agent
         """
         data = self.api.agent_completion(self.project, name, messages)
+        if 'context' in data['message']:
+            return AgentCompletion(data['message']['content'], data['message'].get('context'))
+
         return AgentCompletion(data['message']['content'])
 
     def completion_stream(self, name, messages: List[dict]) -> Iterable[object]: