75 | // eslint-disable-next-line react/prop-types
76 | >(({ className, ...props }, ref) => (
77 | | [role=checkbox]]:translate-y-[2px]',
81 | className
82 | )}
83 | {...props}
84 | />
85 | ))
86 | TableCell.displayName = 'TableCell'
87 |
88 | const TableCaption = React.forwardRef<
89 | HTMLTableCaptionElement,
90 | React.HTMLAttributes
91 | >(({ className, ...props }, ref) => (
92 |
93 | ))
94 | TableCaption.displayName = 'TableCaption'
95 |
96 | export { Table, TableHeader, TableBody, TableFooter, TableHead, TableRow, TableCell, TableCaption }
97 |
--------------------------------------------------------------------------------
/lightrag_webui/src/components/ui/Tabs.tsx:
--------------------------------------------------------------------------------
1 | import * as React from 'react'
2 | import * as TabsPrimitive from '@radix-ui/react-tabs'
3 |
4 | import { cn } from '@/lib/utils'
5 |
6 | const Tabs = TabsPrimitive.Root
7 |
8 | const TabsList = React.forwardRef<
9 | React.ComponentRef,
10 | React.ComponentPropsWithoutRef
11 | >(({ className, ...props }, ref) => (
12 |
20 | ))
21 | TabsList.displayName = TabsPrimitive.List.displayName
22 |
23 | const TabsTrigger = React.forwardRef<
24 | React.ComponentRef,
25 | React.ComponentPropsWithoutRef
26 | >(({ className, ...props }, ref) => (
27 |
35 | ))
36 | TabsTrigger.displayName = TabsPrimitive.Trigger.displayName
37 |
38 | const TabsContent = React.forwardRef<
39 | React.ComponentRef,
40 | React.ComponentPropsWithoutRef
41 | >(({ className, ...props }, ref) => (
42 |
54 | ))
55 | TabsContent.displayName = TabsPrimitive.Content.displayName
56 |
57 | export { Tabs, TabsList, TabsTrigger, TabsContent }
58 |
--------------------------------------------------------------------------------
/lightrag_webui/src/components/ui/Text.tsx:
--------------------------------------------------------------------------------
1 | import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '@/components/ui/Tooltip'
2 | import { cn } from '@/lib/utils'
3 |
4 | const Text = ({
5 | text,
6 | className,
7 | tooltipClassName,
8 | tooltip,
9 | side,
10 | onClick
11 | }: {
12 | text: string
13 | className?: string
14 | tooltipClassName?: string
15 | tooltip?: string
16 | side?: 'top' | 'right' | 'bottom' | 'left'
17 | onClick?: () => void
18 | }) => {
19 | if (!tooltip) {
20 | return (
21 |
27 | )
28 | }
29 |
30 | return (
31 |
32 |
33 |
34 |
40 |
41 |
42 | {tooltip}
43 |
44 |
45 |
46 | )
47 | }
48 |
49 | export default Text
50 |
--------------------------------------------------------------------------------
/lightrag_webui/src/components/ui/Tooltip.tsx:
--------------------------------------------------------------------------------
1 | import * as React from 'react'
2 | import * as TooltipPrimitive from '@radix-ui/react-tooltip'
3 | import { cn } from '@/lib/utils'
4 |
5 | const TooltipProvider = TooltipPrimitive.Provider
6 |
7 | const Tooltip = TooltipPrimitive.Root
8 |
9 | const TooltipTrigger = TooltipPrimitive.Trigger
10 |
11 | const processTooltipContent = (content: string) => {
12 | if (typeof content !== 'string') return content
13 | return (
14 |
15 | {content}
16 |
17 | )
18 | }
19 |
20 | const TooltipContent = React.forwardRef<
21 | React.ComponentRef,
22 | React.ComponentPropsWithoutRef & {
23 | side?: 'top' | 'right' | 'bottom' | 'left'
24 | align?: 'start' | 'center' | 'end'
25 | }
26 | >(({ className, side = 'left', align = 'start', children, ...props }, ref) => {
27 | const contentRef = React.useRef(null);
28 |
29 | React.useEffect(() => {
30 | if (contentRef.current) {
31 | contentRef.current.scrollTop = 0;
32 | }
33 | }, [children]);
34 |
35 | return (
36 |
46 | {typeof children === 'string' ? processTooltipContent(children) : children}
47 |
48 | );
49 | })
50 | TooltipContent.displayName = TooltipPrimitive.Content.displayName
51 |
52 | export { Tooltip, TooltipTrigger, TooltipContent, TooltipProvider }
53 |
--------------------------------------------------------------------------------
/lightrag_webui/src/contexts/TabVisibilityProvider.tsx:
--------------------------------------------------------------------------------
1 | import React, { useState, useEffect, useMemo } from 'react';
2 | import { TabVisibilityContext } from './context';
3 | import { TabVisibilityContextType } from './types';
4 | import { useSettingsStore } from '@/stores/settings';
5 |
6 | interface TabVisibilityProviderProps {
7 | children: React.ReactNode;
8 | }
9 |
10 | /**
11 | * Provider component for the TabVisibility context
12 | * Manages the visibility state of tabs throughout the application
13 | */
14 | export const TabVisibilityProvider: React.FC = ({ children }) => {
15 | // Get current tab from settings store
16 | const currentTab = useSettingsStore.use.currentTab();
17 |
18 | // Initialize visibility state with all tabs visible
19 | const [visibleTabs, setVisibleTabs] = useState>(() => ({
20 | 'documents': true,
21 | 'knowledge-graph': true,
22 | 'retrieval': true,
23 | 'api': true
24 | }));
25 |
26 | // Keep all tabs visible because we use CSS to control TAB visibility instead of React
27 | useEffect(() => {
28 | setVisibleTabs((prev) => ({
29 | ...prev,
30 | 'documents': true,
31 | 'knowledge-graph': true,
32 | 'retrieval': true,
33 | 'api': true
34 | }));
35 | }, [currentTab]);
36 |
37 | // Create the context value with memoization to prevent unnecessary re-renders
38 | const contextValue = useMemo(
39 | () => ({
40 | visibleTabs,
41 | setTabVisibility: (tabId: string, isVisible: boolean) => {
42 | setVisibleTabs((prev) => ({
43 | ...prev,
44 | [tabId]: isVisible,
45 | }));
46 | },
47 | isTabVisible: (tabId: string) => !!visibleTabs[tabId],
48 | }),
49 | [visibleTabs]
50 | );
51 |
52 | return (
53 |
54 | {children}
55 |
56 | );
57 | };
58 |
59 | export default TabVisibilityProvider;
60 |
--------------------------------------------------------------------------------
/lightrag_webui/src/contexts/context.ts:
--------------------------------------------------------------------------------
1 | import { createContext } from 'react';
2 | import { TabVisibilityContextType } from './types';
3 |
4 | // Default context value
5 | const defaultContext: TabVisibilityContextType = {
6 | visibleTabs: {},
7 | setTabVisibility: () => {},
8 | isTabVisible: () => false,
9 | };
10 |
11 | // Create the context
12 | export const TabVisibilityContext = createContext(defaultContext);
13 |
--------------------------------------------------------------------------------
/lightrag_webui/src/contexts/types.ts:
--------------------------------------------------------------------------------
1 | export interface TabVisibilityContextType {
2 | visibleTabs: Record;
3 | setTabVisibility: (tabId: string, isVisible: boolean) => void;
4 | isTabVisible: (tabId: string) => boolean;
5 | }
6 |
--------------------------------------------------------------------------------
/lightrag_webui/src/contexts/useTabVisibility.ts:
--------------------------------------------------------------------------------
1 | import { useContext } from 'react';
2 | import { TabVisibilityContext } from './context';
3 | import { TabVisibilityContextType } from './types';
4 |
5 | /**
6 | * Custom hook to access the tab visibility context
7 | * @returns The tab visibility context
8 | */
9 | export const useTabVisibility = (): TabVisibilityContextType => {
10 | const context = useContext(TabVisibilityContext);
11 |
12 | if (!context) {
13 | throw new Error('useTabVisibility must be used within a TabVisibilityProvider');
14 | }
15 |
16 | return context;
17 | };
18 |
--------------------------------------------------------------------------------
/lightrag_webui/src/features/ApiSite.tsx:
--------------------------------------------------------------------------------
1 | import { useState, useEffect } from 'react'
2 | import { useTabVisibility } from '@/contexts/useTabVisibility'
3 | import { backendBaseUrl } from '@/lib/constants'
4 | import { useTranslation } from 'react-i18next'
5 |
6 | export default function ApiSite() {
7 | const { t } = useTranslation()
8 | const { isTabVisible } = useTabVisibility()
9 | const isApiTabVisible = isTabVisible('api')
10 | const [iframeLoaded, setIframeLoaded] = useState(false)
11 |
12 | // Load the iframe once on component mount
13 | useEffect(() => {
14 | if (!iframeLoaded) {
15 | setIframeLoaded(true)
16 | }
17 | }, [iframeLoaded])
18 |
19 | // Use CSS to hide content when tab is not visible
20 | return (
21 |
22 | {iframeLoaded ? (
23 |
30 | ) : (
31 |
32 |
33 |
34 | {t('apiSite.loading')}
35 |
36 |
37 | )}
38 | = S extends { getState: () => infer T }
56 | ? S & { use: { [K in keyof T]: () => T[K] } }
57 | : never
58 |
59 | export const createSelectors = >>(_store: S) => {
60 | const store = _store as WithSelectors
61 | store.use = {}
62 | for (const k of Object.keys(store.getState())) {
63 | ;(store.use as any)[k] = () => store((s) => s[k as keyof typeof s])
64 | }
65 |
66 | return store
67 | }
68 |
--------------------------------------------------------------------------------
/lightrag_webui/src/main.tsx:
--------------------------------------------------------------------------------
1 | import { StrictMode } from 'react'
2 | import { createRoot } from 'react-dom/client'
3 | import './index.css'
4 | import AppRouter from './AppRouter'
5 | import './i18n.ts';
6 |
7 |
8 |
9 | createRoot(document.getElementById('root')!).render(
10 |
11 |
12 |
13 | )
14 |
--------------------------------------------------------------------------------
/lightrag_webui/src/services/navigation.ts:
--------------------------------------------------------------------------------
1 | import { NavigateFunction } from 'react-router-dom';
2 | import { useAuthStore, useBackendState } from '@/stores/state';
3 | import { useGraphStore } from '@/stores/graph';
4 | import { useSettingsStore } from '@/stores/settings';
5 |
6 | class NavigationService {
7 | private navigate: NavigateFunction | null = null;
8 |
9 | setNavigate(navigate: NavigateFunction) {
10 | this.navigate = navigate;
11 | }
12 |
13 | /**
14 | * Reset all application state to ensure a clean environment.
15 | * This function should be called when:
16 | * 1. User logs out
17 | * 2. Authentication token expires
18 | * 3. Direct access to login page
19 | */
20 | resetAllApplicationState() {
21 | console.log('Resetting all application state...');
22 |
23 | // Reset graph state
24 | const graphStore = useGraphStore.getState();
25 | const sigma = graphStore.sigmaInstance;
26 | graphStore.reset();
27 | graphStore.setGraphDataFetchAttempted(false);
28 | graphStore.setLabelsFetchAttempted(false);
29 | graphStore.setSigmaInstance(null);
30 | graphStore.setIsFetching(false); // Reset isFetching state to prevent data loading issues
31 |
32 | // Reset backend state
33 | useBackendState.getState().clear();
34 |
35 | // Reset retrieval history message while preserving other user preferences
36 | useSettingsStore.getState().setRetrievalHistory([]);
37 |
38 | // Clear authentication state
39 | sessionStorage.clear();
40 |
41 | if (sigma) {
42 | sigma.getGraph().clear();
43 | sigma.kill();
44 | useGraphStore.getState().setSigmaInstance(null);
45 | }
46 | }
47 |
48 | /**
49 | * Handle direct access to login page
50 | * @returns true if it's a direct access, false if navigated from another page
51 | */
52 | handleDirectLoginAccess() {
53 | const isDirectAccess = !document.referrer;
54 | if (isDirectAccess) {
55 | this.resetAllApplicationState();
56 | }
57 | return isDirectAccess;
58 | }
59 |
60 | /**
61 | * Navigate to login page and reset application state
62 | * @param skipReset whether to skip state reset (used for direct access scenario where reset is already handled)
63 | */
64 | navigateToLogin() {
65 | if (!this.navigate) {
66 | console.error('Navigation function not set');
67 | return;
68 | }
69 |
70 | this.resetAllApplicationState();
71 | useAuthStore.getState().logout();
72 |
73 | this.navigate('/login');
74 | }
75 |
76 | navigateToHome() {
77 | if (!this.navigate) {
78 | console.error('Navigation function not set');
79 | return;
80 | }
81 |
82 | this.navigate('/');
83 | }
84 | }
85 |
86 | export const navigationService = new NavigationService();
87 |
--------------------------------------------------------------------------------
/lightrag_webui/src/vite-env.d.ts:
--------------------------------------------------------------------------------
1 | ///
2 |
3 | interface ImportMetaEnv {
4 | readonly VITE_API_PROXY: string
5 | readonly VITE_API_ENDPOINTS: string
6 | readonly VITE_BACKEND_URL: string
7 | }
8 |
9 | interface ImportMeta {
10 | readonly env: ImportMetaEnv
11 | }
12 |
--------------------------------------------------------------------------------
/lightrag_webui/tsconfig.json:
--------------------------------------------------------------------------------
1 | {
2 | "compilerOptions": {
3 | "target": "ES2020",
4 | "useDefineForClassFields": true,
5 | "lib": ["ES2020", "DOM", "DOM.Iterable"],
6 | "module": "ESNext",
7 | "skipLibCheck": true,
8 |
9 | /* Bundler mode */
10 | "moduleResolution": "bundler",
11 | "allowImportingTsExtensions": true,
12 | "isolatedModules": true,
13 | "moduleDetection": "force",
14 | "noEmit": true,
15 | "jsx": "react-jsx",
16 |
17 | /* Linting */
18 | "strict": true,
19 | "noUnusedLocals": true,
20 | "noUnusedParameters": true,
21 | "noFallthroughCasesInSwitch": true,
22 |
23 | /* Paths */
24 | "baseUrl": ".",
25 | "paths": {
26 | "@/*": ["./src/*"]
27 | }
28 | },
29 | "include": ["src", "vite.config.ts", "src/vite-env.d.ts"]
30 | }
31 |
--------------------------------------------------------------------------------
/lightrag_webui/vite.config.ts:
--------------------------------------------------------------------------------
1 | import { defineConfig } from 'vite'
2 | import path from 'path'
3 | import { webuiPrefix } from '@/lib/constants'
4 | import react from '@vitejs/plugin-react-swc'
5 | import tailwindcss from '@tailwindcss/vite'
6 |
7 | // https://vite.dev/config/
8 | export default defineConfig({
9 | plugins: [react(), tailwindcss()],
10 | resolve: {
11 | alias: {
12 | '@': path.resolve(__dirname, './src')
13 | }
14 | },
15 | // base: import.meta.env.VITE_BASE_URL || '/webui/',
16 | base: webuiPrefix,
17 | build: {
18 | outDir: path.resolve(__dirname, '../lightrag/api/webui'),
19 | emptyOutDir: true,
20 | rollupOptions: {
21 | output: {
22 | // Manual chunking strategy
23 | manualChunks: {
24 | // Group React-related libraries into one chunk
25 | 'react-vendor': ['react', 'react-dom', 'react-router-dom'],
26 | // Group graph visualization libraries into one chunk
27 | 'graph-vendor': ['sigma', 'graphology', '@react-sigma/core'],
28 | // Group UI component libraries into one chunk
29 | 'ui-vendor': ['@radix-ui/react-dialog', '@radix-ui/react-popover', '@radix-ui/react-select', '@radix-ui/react-tabs'],
30 | // Group utility libraries into one chunk
31 | 'utils-vendor': ['axios', 'i18next', 'zustand', 'clsx', 'tailwind-merge'],
32 | // Separate feature modules
33 | 'feature-graph': ['./src/features/GraphViewer'],
34 | 'feature-documents': ['./src/features/DocumentManager'],
35 | 'feature-retrieval': ['./src/features/RetrievalTesting'],
36 |
37 | // Mermaid-related modules
38 | 'mermaid-vendor': ['mermaid'],
39 |
40 | // Markdown-related modules
41 | 'markdown-vendor': [
42 | 'react-markdown',
43 | 'rehype-react',
44 | 'remark-gfm',
45 | 'remark-math',
46 | 'react-syntax-highlighter'
47 | ]
48 | },
49 | // Ensure consistent chunk naming format
50 | chunkFileNames: 'assets/[name]-[hash].js',
51 | // Entry file naming format
52 | entryFileNames: 'assets/[name]-[hash].js',
53 | // Asset file naming format
54 | assetFileNames: 'assets/[name]-[hash].[ext]'
55 | }
56 | }
57 | },
58 | server: {
59 | proxy: import.meta.env.VITE_API_PROXY === 'true' && import.meta.env.VITE_API_ENDPOINTS ?
60 | Object.fromEntries(
61 | import.meta.env.VITE_API_ENDPOINTS.split(',').map(endpoint => [
62 | endpoint,
63 | {
64 | target: import.meta.env.VITE_BACKEND_URL || 'http://localhost:9621',
65 | changeOrigin: true,
66 | rewrite: endpoint === '/api' ?
67 | (path) => path.replace(/^\/api/, '') :
68 | endpoint === '/docs' || endpoint === '/openapi.json' ?
69 | (path) => path : undefined
70 | }
71 | ])
72 | ) : {}
73 | }
74 | })
75 |
--------------------------------------------------------------------------------
/reproduce/Step_0.py:
--------------------------------------------------------------------------------
1 | import os
2 | import json
3 | import glob
4 | import argparse
5 |
6 |
7 | def extract_unique_contexts(input_directory, output_directory):
8 | os.makedirs(output_directory, exist_ok=True)
9 |
10 | jsonl_files = glob.glob(os.path.join(input_directory, "*.jsonl"))
11 | print(f"Found {len(jsonl_files)} JSONL files.")
12 |
13 | for file_path in jsonl_files:
14 | filename = os.path.basename(file_path)
15 | name, ext = os.path.splitext(filename)
16 | output_filename = f"{name}_unique_contexts.json"
17 | output_path = os.path.join(output_directory, output_filename)
18 |
19 | unique_contexts_dict = {}
20 |
21 | print(f"Processing file: {filename}")
22 |
23 | try:
24 | with open(file_path, "r", encoding="utf-8") as infile:
25 | for line_number, line in enumerate(infile, start=1):
26 | line = line.strip()
27 | if not line:
28 | continue
29 | try:
30 | json_obj = json.loads(line)
31 | context = json_obj.get("context")
32 | if context and context not in unique_contexts_dict:
33 | unique_contexts_dict[context] = None
34 | except json.JSONDecodeError as e:
35 | print(
36 | f"JSON decoding error in file {filename} at line {line_number}: {e}"
37 | )
38 | except FileNotFoundError:
39 | print(f"File not found: {filename}")
40 | continue
41 | except Exception as e:
42 | print(f"An error occurred while processing file {filename}: {e}")
43 | continue
44 |
45 | unique_contexts_list = list(unique_contexts_dict.keys())
46 | print(
47 | f"There are {len(unique_contexts_list)} unique `context` entries in the file {filename}."
48 | )
49 |
50 | try:
51 | with open(output_path, "w", encoding="utf-8") as outfile:
52 | json.dump(unique_contexts_list, outfile, ensure_ascii=False, indent=4)
53 | print(f"Unique `context` entries have been saved to: {output_filename}")
54 | except Exception as e:
55 | print(f"An error occurred while saving to the file {output_filename}: {e}")
56 |
57 | print("All files have been processed.")
58 |
59 |
60 | if __name__ == "__main__":
61 | parser = argparse.ArgumentParser()
62 | parser.add_argument("-i", "--input_dir", type=str, default="../datasets")
63 | parser.add_argument(
64 | "-o", "--output_dir", type=str, default="../datasets/unique_contexts"
65 | )
66 |
67 | args = parser.parse_args()
68 |
69 | extract_unique_contexts(args.input_dir, args.output_dir)
70 |
--------------------------------------------------------------------------------
/reproduce/Step_1.py:
--------------------------------------------------------------------------------
1 | import os
2 | import json
3 | import time
4 | import asyncio
5 |
6 | from lightrag import LightRAG
7 | from lightrag.kg.shared_storage import initialize_pipeline_status
8 |
9 |
10 | def insert_text(rag, file_path):
11 | with open(file_path, mode="r") as f:
12 | unique_contexts = json.load(f)
13 |
14 | retries = 0
15 | max_retries = 3
16 | while retries < max_retries:
17 | try:
18 | rag.insert(unique_contexts)
19 | break
20 | except Exception as e:
21 | retries += 1
22 | print(f"Insertion failed, retrying ({retries}/{max_retries}), error: {e}")
23 | time.sleep(10)
24 | if retries == max_retries:
25 | print("Insertion failed after exceeding the maximum number of retries")
26 |
27 |
28 | cls = "agriculture"
29 | WORKING_DIR = f"../{cls}"
30 |
31 | if not os.path.exists(WORKING_DIR):
32 | os.mkdir(WORKING_DIR)
33 |
34 |
35 | async def initialize_rag():
36 | rag = LightRAG(working_dir=WORKING_DIR)
37 |
38 | await rag.initialize_storages()
39 | await initialize_pipeline_status()
40 |
41 | return rag
42 |
43 |
44 | def main():
45 | # Initialize RAG instance
46 | rag = asyncio.run(initialize_rag())
47 | insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json")
48 |
49 |
50 | if __name__ == "__main__":
51 | main()
52 |
--------------------------------------------------------------------------------
/reproduce/Step_1_openai_compatible.py:
--------------------------------------------------------------------------------
1 | import os
2 | import json
3 | import time
4 | import asyncio
5 | import numpy as np
6 |
7 | from lightrag import LightRAG
8 | from lightrag.utils import EmbeddingFunc
9 | from lightrag.llm.openai import openai_complete_if_cache, openai_embed
10 | from lightrag.kg.shared_storage import initialize_pipeline_status
11 |
12 |
13 | ## For Upstage API
14 | # please check if embedding_dim=4096 in lightrag.py and llm.py in lightrag direcotry
15 | async def llm_model_func(
16 | prompt, system_prompt=None, history_messages=[], **kwargs
17 | ) -> str:
18 | return await openai_complete_if_cache(
19 | "solar-mini",
20 | prompt,
21 | system_prompt=system_prompt,
22 | history_messages=history_messages,
23 | api_key=os.getenv("UPSTAGE_API_KEY"),
24 | base_url="https://api.upstage.ai/v1/solar",
25 | **kwargs,
26 | )
27 |
28 |
29 | async def embedding_func(texts: list[str]) -> np.ndarray:
30 | return await openai_embed(
31 | texts,
32 | model="solar-embedding-1-large-query",
33 | api_key=os.getenv("UPSTAGE_API_KEY"),
34 | base_url="https://api.upstage.ai/v1/solar",
35 | )
36 |
37 |
38 | ## /For Upstage API
39 |
40 |
41 | def insert_text(rag, file_path):
42 | with open(file_path, mode="r") as f:
43 | unique_contexts = json.load(f)
44 |
45 | retries = 0
46 | max_retries = 3
47 | while retries < max_retries:
48 | try:
49 | rag.insert(unique_contexts)
50 | break
51 | except Exception as e:
52 | retries += 1
53 | print(f"Insertion failed, retrying ({retries}/{max_retries}), error: {e}")
54 | time.sleep(10)
55 | if retries == max_retries:
56 | print("Insertion failed after exceeding the maximum number of retries")
57 |
58 |
59 | cls = "mix"
60 | WORKING_DIR = f"../{cls}"
61 |
62 | if not os.path.exists(WORKING_DIR):
63 | os.mkdir(WORKING_DIR)
64 |
65 |
66 | async def initialize_rag():
67 | rag = LightRAG(
68 | working_dir=WORKING_DIR,
69 | llm_model_func=llm_model_func,
70 | embedding_func=EmbeddingFunc(
71 | embedding_dim=4096, max_token_size=8192, func=embedding_func
72 | ),
73 | )
74 |
75 | await rag.initialize_storages()
76 | await initialize_pipeline_status()
77 |
78 | return rag
79 |
80 |
81 | def main():
82 | # Initialize RAG instance
83 | rag = asyncio.run(initialize_rag())
84 | insert_text(rag, f"../datasets/unique_contexts/{cls}_unique_contexts.json")
85 |
86 |
87 | if __name__ == "__main__":
88 | main()
89 |
--------------------------------------------------------------------------------
/reproduce/Step_2.py:
--------------------------------------------------------------------------------
1 | import json
2 | from openai import OpenAI
3 | from transformers import GPT2Tokenizer
4 |
5 |
6 | def openai_complete_if_cache(
7 | model="gpt-4o", prompt=None, system_prompt=None, history_messages=[], **kwargs
8 | ) -> str:
9 | openai_client = OpenAI()
10 |
11 | messages = []
12 | if system_prompt:
13 | messages.append({"role": "system", "content": system_prompt})
14 | messages.extend(history_messages)
15 | messages.append({"role": "user", "content": prompt})
16 |
17 | response = openai_client.chat.completions.create(
18 | model=model, messages=messages, **kwargs
19 | )
20 | return response.choices[0].message.content
21 |
22 |
23 | tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
24 |
25 |
26 | def get_summary(context, tot_tokens=2000):
27 | tokens = tokenizer.tokenize(context)
28 | half_tokens = tot_tokens // 2
29 |
30 | start_tokens = tokens[1000 : 1000 + half_tokens]
31 | end_tokens = tokens[-(1000 + half_tokens) : 1000]
32 |
33 | summary_tokens = start_tokens + end_tokens
34 | summary = tokenizer.convert_tokens_to_string(summary_tokens)
35 |
36 | return summary
37 |
38 |
39 | clses = ["agriculture"]
40 | for cls in clses:
41 | with open(f"../datasets/unique_contexts/{cls}_unique_contexts.json", mode="r") as f:
42 | unique_contexts = json.load(f)
43 |
44 | summaries = [get_summary(context) for context in unique_contexts]
45 |
46 | total_description = "\n\n".join(summaries)
47 |
48 | prompt = f"""
49 | Given the following description of a dataset:
50 |
51 | {total_description}
52 |
53 | Please identify 5 potential users who would engage with this dataset. For each user, list 5 tasks they would perform with this dataset. Then, for each (user, task) combination, generate 5 questions that require a high-level understanding of the entire dataset.
54 |
55 | Output the results in the following structure:
56 | - User 1: [user description]
57 | - Task 1: [task description]
58 | - Question 1:
59 | - Question 2:
60 | - Question 3:
61 | - Question 4:
62 | - Question 5:
63 | - Task 2: [task description]
64 | ...
65 | - Task 5: [task description]
66 | - User 2: [user description]
67 | ...
68 | - User 5: [user description]
69 | ...
70 | """
71 |
72 | result = openai_complete_if_cache(model="gpt-4o", prompt=prompt)
73 |
74 | file_path = f"../datasets/questions/{cls}_questions.txt"
75 | with open(file_path, "w") as file:
76 | file.write(result)
77 |
78 | print(f"{cls}_questions written to {file_path}")
79 |
--------------------------------------------------------------------------------
/reproduce/Step_3.py:
--------------------------------------------------------------------------------
1 | import re
2 | import json
3 | from lightrag import LightRAG, QueryParam
4 | from lightrag.utils import always_get_an_event_loop
5 |
6 |
7 | def extract_queries(file_path):
8 | with open(file_path, "r") as f:
9 | data = f.read()
10 |
11 | data = data.replace("**", "")
12 |
13 | queries = re.findall(r"- Question \d+: (.+)", data)
14 |
15 | return queries
16 |
17 |
18 | async def process_query(query_text, rag_instance, query_param):
19 | try:
20 | result = await rag_instance.aquery(query_text, param=query_param)
21 | return {"query": query_text, "result": result}, None
22 | except Exception as e:
23 | return None, {"query": query_text, "error": str(e)}
24 |
25 |
26 | def run_queries_and_save_to_json(
27 | queries, rag_instance, query_param, output_file, error_file
28 | ):
29 | loop = always_get_an_event_loop()
30 |
31 | with open(output_file, "a", encoding="utf-8") as result_file, open(
32 | error_file, "a", encoding="utf-8"
33 | ) as err_file:
34 | result_file.write("[\n")
35 | first_entry = True
36 |
37 | for query_text in queries:
38 | result, error = loop.run_until_complete(
39 | process_query(query_text, rag_instance, query_param)
40 | )
41 |
42 | if result:
43 | if not first_entry:
44 | result_file.write(",\n")
45 | json.dump(result, result_file, ensure_ascii=False, indent=4)
46 | first_entry = False
47 | elif error:
48 | json.dump(error, err_file, ensure_ascii=False, indent=4)
49 | err_file.write("\n")
50 |
51 | result_file.write("\n]")
52 |
53 |
54 | if __name__ == "__main__":
55 | cls = "agriculture"
56 | mode = "hybrid"
57 | WORKING_DIR = f"../{cls}"
58 |
59 | rag = LightRAG(working_dir=WORKING_DIR)
60 | query_param = QueryParam(mode=mode)
61 |
62 | queries = extract_queries(f"../datasets/questions/{cls}_questions.txt")
63 | run_queries_and_save_to_json(
64 | queries, rag, query_param, f"{cls}_result.json", f"{cls}_errors.json"
65 | )
66 |
--------------------------------------------------------------------------------
/reproduce/Step_3_openai_compatible.py:
--------------------------------------------------------------------------------
1 | import os
2 | import re
3 | import json
4 | from lightrag import LightRAG, QueryParam
5 | from lightrag.llm.openai import openai_complete_if_cache, openai_embed
6 | from lightrag.utils import EmbeddingFunc, always_get_an_event_loop
7 | import numpy as np
8 |
9 |
10 | ## For Upstage API
11 | # please check if embedding_dim=4096 in lightrag.py and llm.py in lightrag direcotry
12 | async def llm_model_func(
13 | prompt, system_prompt=None, history_messages=[], **kwargs
14 | ) -> str:
15 | return await openai_complete_if_cache(
16 | "solar-mini",
17 | prompt,
18 | system_prompt=system_prompt,
19 | history_messages=history_messages,
20 | api_key=os.getenv("UPSTAGE_API_KEY"),
21 | base_url="https://api.upstage.ai/v1/solar",
22 | **kwargs,
23 | )
24 |
25 |
26 | async def embedding_func(texts: list[str]) -> np.ndarray:
27 | return await openai_embed(
28 | texts,
29 | model="solar-embedding-1-large-query",
30 | api_key=os.getenv("UPSTAGE_API_KEY"),
31 | base_url="https://api.upstage.ai/v1/solar",
32 | )
33 |
34 |
35 | ## /For Upstage API
36 |
37 |
38 | def extract_queries(file_path):
39 | with open(file_path, "r") as f:
40 | data = f.read()
41 |
42 | data = data.replace("**", "")
43 |
44 | queries = re.findall(r"- Question \d+: (.+)", data)
45 |
46 | return queries
47 |
48 |
49 | async def process_query(query_text, rag_instance, query_param):
50 | try:
51 | result = await rag_instance.aquery(query_text, param=query_param)
52 | return {"query": query_text, "result": result}, None
53 | except Exception as e:
54 | return None, {"query": query_text, "error": str(e)}
55 |
56 |
57 | def run_queries_and_save_to_json(
58 | queries, rag_instance, query_param, output_file, error_file
59 | ):
60 | loop = always_get_an_event_loop()
61 |
62 | with open(output_file, "a", encoding="utf-8") as result_file, open(
63 | error_file, "a", encoding="utf-8"
64 | ) as err_file:
65 | result_file.write("[\n")
66 | first_entry = True
67 |
68 | for query_text in queries:
69 | result, error = loop.run_until_complete(
70 | process_query(query_text, rag_instance, query_param)
71 | )
72 |
73 | if result:
74 | if not first_entry:
75 | result_file.write(",\n")
76 | json.dump(result, result_file, ensure_ascii=False, indent=4)
77 | first_entry = False
78 | elif error:
79 | json.dump(error, err_file, ensure_ascii=False, indent=4)
80 | err_file.write("\n")
81 |
82 | result_file.write("\n]")
83 |
84 |
85 | if __name__ == "__main__":
86 | cls = "mix"
87 | mode = "hybrid"
88 | WORKING_DIR = f"../{cls}"
89 |
90 | rag = LightRAG(working_dir=WORKING_DIR)
91 | rag = LightRAG(
92 | working_dir=WORKING_DIR,
93 | llm_model_func=llm_model_func,
94 | embedding_func=EmbeddingFunc(
95 | embedding_dim=4096, max_token_size=8192, func=embedding_func
96 | ),
97 | )
98 | query_param = QueryParam(mode=mode)
99 |
100 | base_dir = "../datasets/questions"
101 | queries = extract_queries(f"{base_dir}/{cls}_questions.txt")
102 | run_queries_and_save_to_json(
103 | queries, rag, query_param, f"{base_dir}/result.json", f"{base_dir}/errors.json"
104 | )
105 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | aiohttp
2 | configparser
3 | future
4 |
5 | # Additional Packages for export Functionality
6 | pandas>=2.0.0
7 |
8 | # Extra libraries are installed when needed using pipmaster
9 |
10 | pipmaster
11 | pydantic
12 | python-dotenv
13 |
14 | # Unicode Collation Algorithm for proper Chinese sorting
15 | pyuca
16 |
17 | setuptools
18 | tenacity
19 |
20 | # LLM packages
21 | tiktoken
22 | xlsxwriter>=3.1.0
23 |
--------------------------------------------------------------------------------
|