|
| 1 | +import mindsdb_sdk |
| 2 | +from uuid import uuid4 |
| 3 | +import os |
| 4 | + |
| 5 | +con = mindsdb_sdk.connect() |
| 6 | + |
| 7 | +open_ai_key = os.getenv('OPENAI_API_KEY') |
| 8 | +model_name = 'gpt-4' |
| 9 | + |
| 10 | +# Now create an agent that will use the model we just created. |
| 11 | +agent = con.agents.create(name=f'mindsdb_sql_agent_{model_name}_{uuid4().hex}', |
| 12 | + model='gpt-4') |
| 13 | + |
| 14 | + |
| 15 | +# Set up a Postgres data source with our new agent. |
| 16 | +data_source = 'postgres' |
| 17 | +connection_args = { |
| 18 | + "user": "demo_user", |
| 19 | + "password": "demo_password", |
| 20 | + "host": "samples.mindsdb.com", |
| 21 | + "port": "5432", |
| 22 | + "database": "demo", |
| 23 | + "schema": "demo_data" |
| 24 | +} |
| 25 | +description = 'mindsdb demo database' |
| 26 | +database = con.databases.create( |
| 27 | + f'mindsdb_sql_agent_datasource_{uuid4().hex}', |
| 28 | + data_source, |
| 29 | + connection_args |
| 30 | +) |
| 31 | + |
| 32 | +# Actually connect the agent to the datasource. |
| 33 | +agent.add_database(database.name, [], description) |
| 34 | + |
| 35 | + |
| 36 | +question = 'How many three-bedroom houses were sold in 2008?' |
| 37 | + |
| 38 | +completion_stream = agent.completion_stream([{'question': question, 'answer': None}]) |
| 39 | + |
| 40 | +# Process the streaming response |
| 41 | +full_response = "" |
| 42 | +for chunk in completion_stream: |
| 43 | + print(chunk) # Print the entire chunk for debugging |
| 44 | + if isinstance(chunk, dict): |
| 45 | + if 'output' in chunk: |
| 46 | + full_response += chunk['output'] |
| 47 | + elif isinstance(chunk, str): |
| 48 | + full_response += chunk |
| 49 | + |
| 50 | +print("\n\nFull response:") |
| 51 | +print(full_response) |
| 52 | + |
| 53 | +con.databases.drop(database.name) |
| 54 | +con.agents.drop(agent.name) |
0 commit comments