Develop support for external database as datasource (#42)

Author: vegito22

apps/consumers.py

@@ -66,9 +66,9 @@ async def _respond_to_event(self, text_data):
                 output_stream = await AppViewSet().run_app_internal_async(self.app_id, self._session_id, request_uuid, request, self.preview)
                 async for output in output_stream:
                     if 'errors' in output or 'session' in output:
-                        await self.send(text_data=output)
+                        await self.send(text_data=json.dumps(output))
                     else:
-                        await self.send(text_data="{\"output\":" + output + '}')
+                        await self.send(text_data=json.dumps({'output': output}))
 
                 await self.send(text_data=json.dumps({'event': 'done'}))
             except Exception as e:

apps/handlers/app_runnner.py

@@ -219,9 +219,9 @@ async def stream_output():
                             await asyncio.sleep(0.0001)
                             if not metadata_sent:
                                 metadata_sent = True
-                                yield json.dumps({'session': {'id': app_session['uuid']}, 'csp': csp, 'template': template}) + '\n'
+                                yield {'session': {'id': app_session['uuid']}, 'csp': csp, 'template': template}
                             output = next(output_iter)
-                            yield json.dumps(output) + '\n'
+                            yield output
                     except StopIteration:
                         pass
                     except Exception as e:

apps/tasks.py

@@ -114,6 +114,8 @@ def resync_data_entry_task(datasource: DataSource, entry_data: DataSourceEntry):
 
 def delete_data_source_task(datasource):
     datasource_type = datasource.type
+    if datasource_type.is_external_datasource:
+        return
     datasource_entry_handler_cls = DataSourceTypeFactory.get_datasource_type_handler(
         datasource_type,
     )

client/src/components/datasource/AddDataSourceModal.jsx

@@ -136,17 +136,21 @@ export function AddDataSourceModal({
                 .post("/api/datasources", {
                   name: dataSourceName,
                   type: dataSourceType.id,
+                  config: dataSourceType.is_external_datasource ? formData : {},
                 })
                 .then((response) => {
-                  const dataSource = response.data;
-                  setDataSources([...dataSources, dataSource]);
-                  axios()
-                    .post(`/api/datasources/${dataSource.uuid}/add_entry`, {
-                      entry_data: formData,
-                    })
-                    .then((response) => {
-                      dataSourceAddedCb(dataSource);
-                    });
+                  // External data sources do not support adding entries
+                  if (!dataSourceType.is_external_datasource) {
+                    const dataSource = response.data;
+                    setDataSources([...dataSources, dataSource]);
+                    axios()
+                      .post(`/api/datasources/${dataSource.uuid}/add_entry`, {
+                        entry_data: formData,
+                      })
+                      .then((response) => {
+                        dataSourceAddedCb(dataSource);
+                      });
+                  }
                 });
               handleCancelCb();
               enqueueSnackbar(

client/src/pages/data.jsx

@@ -9,13 +9,15 @@ import {
   Chip,
   Grid,
   Stack,
+  Tooltip,
 } from "@mui/material";
 
 import { TextareaAutosize } from "@mui/base";
 
 import DeleteOutlineOutlinedIcon from "@mui/icons-material/DeleteOutlineOutlined";
 import AddOutlinedIcon from "@mui/icons-material/AddOutlined";
 import SyncOutlinedIcon from "@mui/icons-material/SyncOutlined";
+import SettingsEthernetIcon from "@mui/icons-material/SettingsEthernet";
 import PeopleOutlineOutlinedIcon from "@mui/icons-material/PeopleOutlineOutlined";
 import PersonOutlineOutlinedIcon from "@mui/icons-material/PersonOutlineOutlined";
 
@@ -192,51 +194,64 @@ export default function DataPage() {
     {
       title: "Action",
       key: "operation",
-      render: (record) => (
-        <Box>
-          <IconButton
-            disabled={!record.isUserOwned}
-            onClick={() => {
-              setModalTitle("Add New Data Entry");
-              setSelectedDataSource(record);
-              setAddDataSourceModalOpen(true);
-            }}
-          >
-            <AddOutlinedIcon />
-          </IconButton>
-          <IconButton
-            disabled={!record.isUserOwned}
-            onClick={() => {
-              setDeleteId(record);
-              setDeleteModalTitle("Delete Data Source");
-              setDeleteModalMessage(
-                <div>
-                  Are you sure you want to delete{" "}
-                  <span style={{ fontWeight: "bold" }}>{record.name}</span> ?
-                </div>,
-              );
-              setDeleteConfirmationModalOpen(true);
-            }}
-          >
-            <DeleteOutlineOutlinedIcon />
-          </IconButton>
-          {profileFlags.IS_ORGANIZATION_MEMBER && record.isUserOwned && (
+      render: (record) => {
+        return (
+          <Box>
+            {!record?.type?.is_external_datasource && (
+              <IconButton
+                disabled={!record.isUserOwned}
+                onClick={() => {
+                  setModalTitle("Add New Data Entry");
+                  setSelectedDataSource(record);
+                  setAddDataSourceModalOpen(true);
+                }}
+              >
+                <AddOutlinedIcon />
+              </IconButton>
+            )}
+            {record?.type?.is_external_datasource && (
+              <Tooltip title="External Connection">
+                <span>
+                  <IconButton disabled={true}>
+                    <SettingsEthernetIcon />
+                  </IconButton>
+                </span>
+              </Tooltip>
+            )}
             <IconButton
+              disabled={!record.isUserOwned}
               onClick={() => {
-                setModalTitle("Share Datasource");
-                setSelectedDataSource(record);
-                setShareDataSourceModalOpen(true);
+                setDeleteId(record);
+                setDeleteModalTitle("Delete Data Source");
+                setDeleteModalMessage(
+                  <div>
+                    Are you sure you want to delete{" "}
+                    <span style={{ fontWeight: "bold" }}>{record.name}</span> ?
+                  </div>,
+                );
+                setDeleteConfirmationModalOpen(true);
               }}
             >
-              {record.visibility === 0 ? (
-                <PersonOutlineOutlinedIcon />
-              ) : (
-                <PeopleOutlineOutlinedIcon />
-              )}
+              <DeleteOutlineOutlinedIcon />
             </IconButton>
-          )}
-        </Box>
-      ),
+            {profileFlags.IS_ORGANIZATION_MEMBER && record.isUserOwned && (
+              <IconButton
+                onClick={() => {
+                  setModalTitle("Share Datasource");
+                  setSelectedDataSource(record);
+                  setShareDataSourceModalOpen(true);
+                }}
+              >
+                {record.visibility === 0 ? (
+                  <PersonOutlineOutlinedIcon />
+                ) : (
+                  <PeopleOutlineOutlinedIcon />
+                )}
+              </IconButton>
+            )}
+          </Box>
+        );
+      },
     },
   ];
 

common/blocks/data/store/vectorstore/weaviate.py

@@ -80,6 +80,10 @@ class WeaviateConfiguration(BaseModel):
     weaviate_rw_api_key: Optional[str] = None
     embeddings_rate_limit: Optional[int] = 3000
     default_batch_size: Optional[int] = 20
+    username: Optional[str] = None
+    password: Optional[str] = None
+    api_key: Optional[str] = None
+    additional_headers: Optional[dict] = {}
 
 
 class Weaviate(VectorStoreInterface):
@@ -130,8 +134,8 @@ def check_batch_result(results: Optional[List[Dict[str, Any]]]) -> None:
                                 json.dumps(result['result']['errors']),
                             ),
                         )
-
-        headers = {}
+        
+        headers = configuration.additional_headers
         if configuration.openai_key is not None:
             headers['X-OpenAI-Api-Key'] = configuration.openai_key
         if configuration.cohere_api_key is not None:
@@ -144,10 +148,25 @@ def check_batch_result(results: Optional[List[Dict[str, Any]]]) -> None:
             headers['authorization'] = 'Bearer ' + \
                 configuration.weaviate_rw_api_key
 
-        self._client = weaviate.Client(
-            url=configuration.url,
-            additional_headers=headers,
-        )
+        if configuration.username is not None and configuration.password is not None:
+            self._client = weaviate.Client(
+                url=configuration.url,
+                auth_client_secret=weaviate.AuthClientPassword(
+                    username=configuration.username, password=configuration.password),
+                additional_headers=headers,
+            )
+        elif configuration.api_key is not None:
+            self._client = weaviate.Client(
+                url=configuration.url,
+                auth_client_secret=weaviate.AuthApiKey(
+                    api_key=configuration.api_key),
+                additional_headers=headers,
+            )
+        else:
+            self._client = weaviate.Client(
+                url=configuration.url,
+                additional_headers=headers,
+            )
 
         self.client.batch.configure(
             batch_size=DEFAULT_BATCH_SIZE,
@@ -234,6 +253,7 @@ def similarity_search(self, index_name: str, document_query: DocumentQuery, **kw
         properties = [document_query.page_content_key]
         for key in document_query.metadata.get('additional_properties', []):
             properties.append(key)
+        additional_metadata_properties = document_query.metadata.get('metadata_properties', ['id', 'certainty', 'distance'])
 
         if kwargs.get('search_distance'):
             nearText['certainty'] = kwargs.get('search_distance')
@@ -254,7 +274,7 @@ def similarity_search(self, index_name: str, document_query: DocumentQuery, **kw
                 query_obj = query_obj.with_where(whereFilter)
             query_response = query_obj.with_near_text(nearText).with_limit(
                 document_query.limit,
-            ).with_additional(['id', 'certainty', 'distance']).do()
+            ).with_additional(additional_metadata_properties).do()
         except Exception as e:
             logger.error('Error in similarity search: %s' % e)
             raise e

common/utils/models.py

@@ -0,0 +1,38 @@
+from typing import Optional
+
+import orjson as json
+from pydantic import BaseModel
+
+class Config(BaseModel):
+    """
+    Base class for config type models stored in the database. Supports optional encryption.
+    """
+    config_type: str
+    is_encrypted: bool = False
+    data: str = ''
+
+    def to_dict(self, encrypt_fn):
+        return {
+            'config_type': self.config_type,
+            'is_encrypted': self.is_encrypted,
+            'data': self.get_data(encrypt_fn),
+        }
+
+    def from_dict(self, dict_data, decrypt_fn):
+        self.config_type = dict_data.get('config_type')
+        self.is_encrypted = dict_data.get('is_encrypted')
+        self.set_data(dict_data.get('data'), decrypt_fn)
+
+        # Use the data from the dict to populate the fields
+        self.__dict__.update(json.loads(self.data))
+
+        return self.dict(exclude={'is_encrypted', 'config_type', 'data'})
+
+    def get_data(self, encrypt_fn):
+        data = self.json(exclude={'is_encrypted', 'config_type', 'data'})
+        return encrypt_fn(data).decode('utf-8') if self.is_encrypted else data
+
+    def set_data(self, data, decrypt_fn):
+        self.data = data
+        if self.is_encrypted:
+            self.data = decrypt_fn(data)

datasources/apis.py

@@ -117,11 +117,28 @@ def post(self, request):
         datasource_type = get_object_or_404(
             DataSourceType, id=request.data['type'],
         )
-        datasource = DataSource.objects.create(
+        
+        datasource = DataSource(
             name=request.data['name'],
             owner=owner,
             type=datasource_type,
         )
+        # If this is an external data source, then we need to save the config in datasource object
+        if datasource_type.is_external_datasource:
+            datasource_type_cls = DataSourceTypeFactory.get_datasource_type_handler(datasource.type)
+            if not datasource_type_cls:
+                logger.error(
+                    'No handler found for data source type {datasource.type}',
+                )
+                return DRFResponse({'errors': ['No handler found for data source type']}, status=400)
+
+            datasource_handler: DataSourceProcessor = datasource_type_cls(datasource)
+            if not datasource_handler:
+                logger.error(f'Error while creating handler for data source {datasource.name}')
+                return DRFResponse({'errors': ['Error while creating handler for data source type']}, status=400)
+            config = datasource_type_cls.process_validate_config(request.data['config'], datasource)
+            datasource.config = config
+        
         datasource.save()
         return DRFResponse(DataSourceSerializer(instance=datasource).data, status=201)
 
@@ -147,6 +164,9 @@ def add_entry(self, request, uid):
         datasource = get_object_or_404(
             DataSource, uuid=uuid.UUID(uid), owner=request.user,
         )
+        if datasource and datasource.type.is_external_datasource:
+            return DRFResponse({'errors': ['Cannot add entry to external data source']}, status=400)
+        
         entry_data = request.data['entry_data']
         entry_metadata = dict(map(lambda x: (f'md_{x}', request.data['entry_metadata'][x]), request.data['entry_metadata'].keys())) if 'entry_metadata' in request.data else {
         }