Merge branch 'langgenius:main' into main

langgenius · May 22, 2024 · 3182f68 · 3182f68
2 parents a254ac2 + f259278
commit 3182f68
Show file tree

Hide file tree

Showing 5 changed files with 93 additions and 3 deletions.
diff --git a/api/commands.py b/api/commands.py
@@ -448,9 +448,64 @@ def convert_to_agent_apps():
     click.echo(click.style('Congratulations! Converted {} agent apps.'.format(len(proceeded_app_ids)), fg='green'))
 
 
+@click.command('add-qdrant-doc-id-index', help='add qdrant doc_id index.')
+@click.option('--field', default='metadata.doc_id', prompt=False, help='index field , default is metadata.doc_id.')
+def add_qdrant_doc_id_index(field: str):
+    click.echo(click.style('Start add qdrant doc_id index.', fg='green'))
+    config = current_app.config
+    vector_type = config.get('VECTOR_STORE')
+    if vector_type != "qdrant":
+        click.echo(click.style('Sorry, only support qdrant vector store.', fg='red'))
+        return
+    create_count = 0
+
+    try:
+        bindings = db.session.query(DatasetCollectionBinding).all()
+        if not bindings:
+            click.echo(click.style('Sorry, no dataset collection bindings found.', fg='red'))
+            return
+        import qdrant_client
+        from qdrant_client.http.exceptions import UnexpectedResponse
+        from qdrant_client.http.models import PayloadSchemaType
+
+        from core.rag.datasource.vdb.qdrant.qdrant_vector import QdrantConfig
+        for binding in bindings:
+            qdrant_config = QdrantConfig(
+                endpoint=config.get('QDRANT_URL'),
+                api_key=config.get('QDRANT_API_KEY'),
+                root_path=current_app.root_path,
+                timeout=config.get('QDRANT_CLIENT_TIMEOUT'),
+                grpc_port=config.get('QDRANT_GRPC_PORT'),
+                prefer_grpc=config.get('QDRANT_GRPC_ENABLED')
+            )
+            try:
+                client = qdrant_client.QdrantClient(**qdrant_config.to_qdrant_params())
+                # create payload index
+                client.create_payload_index(binding.collection_name, field,
+                                            field_schema=PayloadSchemaType.KEYWORD)
+                create_count += 1
+            except UnexpectedResponse as e:
+                # Collection does not exist, so return
+                if e.status_code == 404:
+                    click.echo(click.style(f'Collection not found, collection_name:{binding.collection_name}.', fg='red'))
+                    continue
+                # Some other error occurred, so re-raise the exception
+                else:
+                    click.echo(click.style(f'Failed to create qdrant index, collection_name:{binding.collection_name}.', fg='red'))
+
+    except Exception as e:
+        click.echo(click.style('Failed to create qdrant client.', fg='red'))
+
+    click.echo(
+        click.style(f'Congratulations! Create {create_count} collection indexes.',
+                    fg='green'))
+
+
 def register_commands(app):
     app.cli.add_command(reset_password)
     app.cli.add_command(reset_email)
     app.cli.add_command(reset_encrypt_key_pair)
     app.cli.add_command(vdb_migrate)
     app.cli.add_command(convert_to_agent_apps)
+    app.cli.add_command(add_qdrant_doc_id_index)
+
diff --git a/web/app/(shareLayout)/webapp-signin/page.tsx b/web/app/(shareLayout)/webapp-signin/page.tsx
@@ -6,7 +6,7 @@ import React, { useEffect, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import Toast from '@/app/components/base/toast'
 import Button from '@/app/components/base/button'
-import { fetchSystemFeatures, fetchWebOIDCSSOUrl, fetchWebSAMLSSOUrl } from '@/service/share'
+import { fetchSystemFeatures, fetchWebOAuth2SSOUrl, fetchWebOIDCSSOUrl, fetchWebSAMLSSOUrl } from '@/service/share'
 import LogoSite from '@/app/components/base/logo/logo-site'
 import { setAccessToken } from '@/app/components/share/utils'
 
@@ -90,6 +90,13 @@ const WebSSOForm: FC = () => {
         setIsLoading(false)
       })
     }
+    else if (protocal === 'oauth2') {
+      fetchWebOAuth2SSOUrl(appCode, redirectUrl).then((res) => {
+        router.push(res.url)
+      }).finally(() => {
+        setIsLoading(false)
+      })
+    }
     else {
       Toast.notify({
         type: 'error',

diff --git a/web/app/signin/userSSOForm.tsx b/web/app/signin/userSSOForm.tsx
@@ -5,7 +5,7 @@ import type { FC } from 'react'
 import { useEffect, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import Toast from '@/app/components/base/toast'
-import { getUserOIDCSSOUrl, getUserSAMLSSOUrl } from '@/service/sso'
+import { getUserOAuth2SSOUrl, getUserOIDCSSOUrl, getUserSAMLSSOUrl } from '@/service/sso'
 import Button from '@/app/components/base/button'
 
 type UserSSOFormProps = {
@@ -47,14 +47,29 @@ const UserSSOForm: FC<UserSSOFormProps> = ({
         setIsLoading(false)
       })
     }
-    else {
+    else if (protocol === 'oidc') {
       getUserOIDCSSOUrl().then((res) => {
         document.cookie = `user-oidc-state=${res.state}`
         router.push(res.url)
       }).finally(() => {
         setIsLoading(false)
       })
     }
+    else if (protocol === 'oauth2') {
+      getUserOAuth2SSOUrl().then((res) => {
+        document.cookie = `user-oauth2-state=${res.state}`
+        router.push(res.url)
+      }).finally(() => {
+        setIsLoading(false)
+      })
+    }
+    else {
+      Toast.notify({
+        type: 'error',
+        message: 'invalid SSO protocol',
+      })
+      setIsLoading(false)
+    }
   }
 
   return (

diff --git a/web/service/share.ts b/web/service/share.ts
@@ -159,6 +159,15 @@ export const fetchWebOIDCSSOUrl = async (appCode: string, redirectUrl: string) =
   }) as Promise<{ url: string }>
 }
 
+export const fetchWebOAuth2SSOUrl = async (appCode: string, redirectUrl: string) => {
+  return (getAction('get', false))(getUrl('/enterprise/sso/oauth2/login', false, ''), {
+    params: {
+      app_code: appCode,
+      redirect_url: redirectUrl,
+    },
+  }) as Promise<{ url: string }>
+}
+
 export const fetchAppMeta = async (isInstalledApp: boolean, installedAppId = '') => {
   return (getAction('get', isInstalledApp))(getUrl('meta', isInstalledApp, installedAppId)) as Promise<AppMeta>
 }

diff --git a/web/service/sso.ts b/web/service/sso.ts
@@ -7,3 +7,7 @@ export const getUserSAMLSSOUrl = () => {
 export const getUserOIDCSSOUrl = () => {
   return get<{ url: string; state: string }>('/enterprise/sso/oidc/login')
 }
+
+export const getUserOAuth2SSOUrl = () => {
+  return get<{ url: string; state: string }>('/enterprise/sso/oauth2/login')
+}