Skip to content

Commit

Permalink
enh: poke's data source view (#2601)
Browse files Browse the repository at this point in the history
* enh: poke's data source view

* nit
  • Loading branch information
spolu authored Nov 20, 2023
1 parent 07ccde5 commit 80ef724
Show file tree
Hide file tree
Showing 5 changed files with 107 additions and 8 deletions.
46 changes: 46 additions & 0 deletions core/bin/dust_api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1028,6 +1028,48 @@ async fn data_sources_register(
"data_source": {
"created": ds.created(),
"data_source_id": ds.data_source_id(),
"qdrant_collection": ds.qdrant_collection(),
"config": ds.config(),
},
})),
}),
),
},
}
}

async fn data_sources_retrieve(
extract::Path((project_id, data_source_id)): extract::Path<(i64, String)>,
extract::Extension(state): extract::Extension<Arc<APIState>>,
) -> (StatusCode, Json<APIResponse>) {
let project = project::Project::new_from_id(project_id);
match state
.store
.load_data_source(&project, &data_source_id)
.await
{
Err(e) => error_response(
StatusCode::INTERNAL_SERVER_ERROR,
"internal_server_error",
"Failed to retrieve data source",
Some(e),
),
Ok(ds) => match ds {
None => error_response(
StatusCode::NOT_FOUND,
"data_source_not_found",
&format!("No data source found for id `{}`", data_source_id),
None,
),
Some(ds) => (
StatusCode::OK,
Json(APIResponse {
error: None,
response: Some(json!({
"data_source": {
"created": ds.created(),
"data_source_id": ds.data_source_id(),
"qdrant_collection": ds.qdrant_collection(),
"config": ds.config(),
},
})),
Expand Down Expand Up @@ -2233,6 +2275,10 @@ fn main() {
"/projects/:project_id/data_sources",
post(data_sources_register),
)
.route(
"/projects/:project_id/data_sources/:data_source_id",
get(data_sources_retrieve),
)
.route(
"/projects/:project_id/data_sources/:data_source_id/documents/:document_id/versions",
get(data_sources_documents_versions_list),
Expand Down
20 changes: 20 additions & 0 deletions front/lib/core_api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ export type CoreAPIDataSourceConfig = {
export type CoreAPIDataSource = {
created: number;
data_source_id: string;
qdrant_collection: string;
config: CoreAPIDataSourceConfig;
};

Expand Down Expand Up @@ -482,6 +483,25 @@ export const CoreAPI = {
return _resultFromResponse(response);
},

async getDataSource({
projectId,
dataSourceId,
}: {
projectId: string;
dataSourceId: string;
}): Promise<CoreAPIResponse<{ data_source: CoreAPIDataSource }>> {
const response = await fetch(
`${CORE_API}/projects/${projectId}/data_sources/${dataSourceId}`,
{
headers: {
"Content-Type": "application/json",
},
}
);

return _resultFromResponse(response);
},

async deleteDataSource({
projectId,
dataSourceName,
Expand Down
9 changes: 7 additions & 2 deletions front/lib/swr.ts
Original file line number Diff line number Diff line change
Expand Up @@ -159,11 +159,16 @@ export function useDocuments(
owner: WorkspaceType,
dataSource: { name: string },
limit: number,
offset: number
offset: number,
asDustSuperUser?: boolean
) {
const documentsFetcher: Fetcher<GetDocumentsResponseBody> = fetcher;
const { data, error } = useSWR(
`/api/w/${owner.sId}/data_sources/${dataSource.name}/documents?limit=${limit}&offset=${offset}`,
`/api/w/${owner.sId}/data_sources/${
dataSource.name
}/documents?limit=${limit}&offset=${offset}${
asDustSuperUser ? "&asDustSuperUser=true" : ""
}`,
documentsFetcher
);

Expand Down
11 changes: 7 additions & 4 deletions front/pages/api/w/[wId]/data_sources/[name]/documents/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,13 @@ async function handler(
res: NextApiResponse<GetDocumentsResponseBody>
): Promise<void> {
const session = await getSession(req, res);
const auth = await Authenticator.fromSession(
session,
req.query.wId as string
);
const auth =
req.query.asDustSuperUser === "true"
? await Authenticator.fromSession(session, req.query.wId as string)
: await Authenticator.fromSuperUserSession(
session,
req.query.wId as string
);

const owner = auth.workspace();
if (!owner) {
Expand Down
29 changes: 27 additions & 2 deletions front/pages/poke/[wId]/data_sources/[name]/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import PokeNavbar from "@app/components/poke/PokeNavbar";
import { getDataSource } from "@app/lib/api/data_sources";
import { Authenticator, getSession } from "@app/lib/auth";
import { ConnectorsAPI, ConnectorType } from "@app/lib/connectors_api";
import { CoreAPI, CoreAPIDataSource } from "@app/lib/core_api";
import { getDisplayNameForDocument } from "@app/lib/data_sources";
import { useDocuments } from "@app/lib/swr";
import { timeAgoFrom } from "@app/lib/utils";
Expand All @@ -23,6 +24,7 @@ import { WorkspaceType } from "@app/types/user";
export const getServerSideProps: GetServerSideProps<{
owner: WorkspaceType;
dataSource: DataSourceType;
coreDataSource: CoreAPIDataSource;
connector: ConnectorType | null;
}> = async (context) => {
const wId = context.params?.wId;
Expand Down Expand Up @@ -66,6 +68,17 @@ export const getServerSideProps: GetServerSideProps<{
};
}

const coreDataSourceRes = await CoreAPI.getDataSource({
projectId: dataSource.dustAPIProjectId,
dataSourceId: dataSource.name,
});

if (coreDataSourceRes.isErr()) {
return {
notFound: true,
};
}

let connector: ConnectorType | null = null;
if (dataSource.connectorId) {
const connectorRes = await ConnectorsAPI.getConnector(
Expand All @@ -80,6 +93,7 @@ export const getServerSideProps: GetServerSideProps<{
props: {
owner,
dataSource,
coreDataSource: coreDataSourceRes.value.data_source,
connector,
},
};
Expand All @@ -88,13 +102,14 @@ export const getServerSideProps: GetServerSideProps<{
const DataSourcePage = ({
owner,
dataSource,
coreDataSource,
connector,
}: InferGetServerSidePropsType<typeof getServerSideProps>) => {
const [limit] = useState(10);
const [offset, setOffset] = useState(0);

const { documents, total, isDocumentsLoading, isDocumentsError } =
useDocuments(owner, dataSource, limit, offset);
useDocuments(owner, dataSource, limit, offset, true);

const [displayNameByDocId, setDisplayNameByDocId] = useState<
Record<string, string>
Expand Down Expand Up @@ -134,9 +149,19 @@ const DataSourcePage = ({

<div className="my-8 flex flex-col gap-y-4">
<JsonViewer value={dataSource} rootName={false} />
<JsonViewer value={coreDataSource} rootName={false} />
<JsonViewer value={connector} rootName={false} />
</div>

<div className="flex flex-row">
<a
href={`https://app.datadoghq.eu/logs?query=service%3Acore%20%22DSSTAT%20Finished%20searching%20Qdrant%20documents%22%20%22${coreDataSource.qdrant_collection}%22%20&cols=host%2Cservice&index=%2A&messageDisplay=inline&refresh_mode=sliding&stream_sort=desc&view=spans&viz=stream&live=true`}
className="text-sm text-blue-500"
>
Datadog: Logs DSSTAT Qdrant search
</a>
</div>

<div className="mt-4 flex flex-row">
<div className="flex flex-1">
<div className="flex flex-col">
Expand Down Expand Up @@ -201,7 +226,7 @@ const DataSourcePage = ({
icon={EyeIcon}
onClick={() => {
window.confirm(
"Are you sure you want to access this sensible user data?"
"Are you sure you want to access this sensible user data? (Access will be logged)"
);
void router.push(
`/poke/${owner.sId}/data_sources/${
Expand Down

0 comments on commit 80ef724

Please sign in to comment.