Merge pull request #2627 from stakwork/feature/ai-v2

feat: show only sources for ai search

src/components/App/SideBar/AiSearch/index.tsx

@@ -4,16 +4,17 @@ import styled from 'styled-components'
 import SearchIcon from '~/components/Icons/SearchIcon'
 import { SearchBar } from '~/components/SearchBar'
 import { Flex } from '~/components/common/Flex'
-import { useHasAiChatsResponseLoading } from '~/stores/useAiSummaryStore'
+import { useAiSummaryStore, useHasAiChatsResponseLoading } from '~/stores/useAiSummaryStore'
 import { useDataStore } from '~/stores/useDataStore'
 import { useUserStore } from '~/stores/useUserStore'
 import { colors } from '~/utils'
 
 export const AiSearch = () => {
   const form = useForm<{ search: string }>({ mode: 'onChange' })
-  const { fetchData, setAbortRequests } = useDataStore((s) => s)
+  const { setAbortRequests } = useDataStore((s) => s)
   const { setBudget } = useUserStore((s) => s)
   const { reset } = form
+  const fetchAIData = useAiSummaryStore((s) => s.fetchAIData)
 
   const isLoading = useHasAiChatsResponseLoading()
 
@@ -22,7 +23,7 @@ export const AiSearch = () => {
       return
     }
 
-    fetchData(setBudget, setAbortRequests, search)
+    fetchAIData(setBudget, setAbortRequests, search)
     reset({ search: '' })
   })
 

src/components/App/SideBar/AiSummary/AiQuestions/index.tsx

@@ -5,7 +5,7 @@ import styled from 'styled-components'
 import PlusIcon from '~/components/Icons/PlusIcon'
 import StackIcon from '~/components/Icons/StackIcon'
 import { Flex } from '~/components/common/Flex'
-import { useHasAiChatsResponseLoading } from '~/stores/useAiSummaryStore'
+import { useAiSummaryStore, useHasAiChatsResponseLoading } from '~/stores/useAiSummaryStore'
 import { useDataStore } from '~/stores/useDataStore'
 import { useUserStore } from '~/stores/useUserStore'
 import { colors } from '~/utils'
@@ -16,17 +16,19 @@ type Props = {
 
 // eslint-disable-next-line no-underscore-dangle
 const _AiQuestions = ({ questions }: Props) => {
-  const { fetchData, setAbortRequests } = useDataStore((s) => s)
+  const { setAbortRequests } = useDataStore((s) => s)
   const [setBudget] = useUserStore((s) => [s.setBudget])
   const hasLoadingResponse = useHasAiChatsResponseLoading()
 
+  const fetchAIData = useAiSummaryStore((s) => s.fetchAIData)
+
   const handleSubmitQuestion = (question: string) => {
     if (hasLoadingResponse) {
       return
     }
 
     if (question) {
-      fetchData(setBudget, setAbortRequests, question)
+      fetchAIData(setBudget, setAbortRequests, question)
     }
   }
 

src/components/App/SideBar/AiSummary/AiSources/index.tsx

@@ -1,13 +1,14 @@
 import { Slide } from '@mui/material'
 import Button from '@mui/material/Button'
-import { memo, useCallback, useRef, useState } from 'react'
+import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import styled from 'styled-components'
 import { Episode } from '~/components/App/SideBar/Relevance/Episode'
 import ChevronDownIcon from '~/components/Icons/ChevronDownIcon'
 import ChevronUpIcon from '~/components/Icons/ChevronUpIcon'
 import SourcesIcon from '~/components/Icons/SourcesIcon'
 import { ScrollView } from '~/components/ScrollView'
 import { Flex } from '~/components/common/Flex'
+import { useAiSummaryStore } from '~/stores/useAiSummaryStore'
 import { useDataStore } from '~/stores/useDataStore'
 import { useUpdateSelectedNode } from '~/stores/useGraphStore'
 import { NodeExtended } from '~/types'
@@ -17,29 +18,69 @@ import { adaptTweetNode } from '~/utils/twitterAdapter'
 
 type Props = {
   sourceIds: string[]
+  question: string
+}
+
+const EDGE = {
+  edge_type: 'POSTED',
+  properties: {
+    date_added_to_graph: '1737561660.0435429',
+    weight: 1,
+  },
+  ref_id: '7efabdc1-b494-4d8c-8d55-5062ce1237d3',
+  source: 'efa4819e-a54e-49dd-858f-2aed5cf10940',
+  target: '0ba6ef37-bf92-4094-89a0-893f05d34e6c',
 }
 
 // eslint-disable-next-line no-underscore-dangle
-const _AiSources = ({ sourceIds }: Props) => {
+const _AiSources = ({ sourceIds, question }: Props) => {
   const scrollViewRef = useRef<HTMLDivElement | null>(null)
   const [showAll, setShowAll] = useState(false)
+  const addNewNode = useDataStore((s) => s.addNewNode)
+
+  const beenAdded = useRef(false)
 
-  const { dataInitial } = useDataStore((s) => s)
+  const { dataInitial } = useAiSummaryStore((s) => s)
   const setSelectedNode = useUpdateSelectedNode()
 
+  const currentNodes = useMemo(
+    () => dataInitial?.nodes.filter((i) => sourceIds.includes(i.ref_id)) || [],
+    [dataInitial?.nodes, sourceIds],
+  )
+
+  useEffect(() => {
+    if (!currentNodes.length || beenAdded.current) {
+      return
+    }
+
+    const edges = currentNodes.map((i, index) => ({
+      ...EDGE,
+      source: question,
+      target: i.ref_id,
+      ref_id: `${String(+new Date())}-${index}`,
+      edge_type: 'IS_SOURCE',
+      properties: {
+        date_added_to_graph: String(new Date()),
+        weight: 1,
+      },
+    }))
+
+    beenAdded.current = true
+
+    addNewNode({ nodes: currentNodes, edges })
+  }, [currentNodes, addNewNode, question])
+
+  const handleLoadMoreClick = () => setShowAll(!showAll)
+
+  const visibleNodes = showAll ? currentNodes : [...currentNodes].slice(0, 3)
+
   const handleNodeClick = useCallback(
     (node: NodeExtended) => {
       setSelectedNode(node)
     },
     [setSelectedNode],
   )
 
-  const handleLoadMoreClick = () => setShowAll(!showAll)
-
-  const currentNodes = dataInitial?.nodes.filter((i) => sourceIds.includes(i.ref_id)) || []
-
-  const visibleNodes = showAll ? currentNodes : [...currentNodes].slice(0, 3)
-
   return (
     <SectionWrapper>
       <Slide direction="right" in mountOnEnter>

src/components/App/SideBar/AiSummary/index.tsx

@@ -140,7 +140,9 @@ export const AiSummary = ({ question, response, refId }: Props) => {
           ) : (
             <AiQuestions questions={response.questions || []} />
           )}
-          {(response?.sources || []).length ? <AiSources sourceIds={response.sources || []} /> : null}
+          {(response?.sources || []).length ? (
+            <AiSources question={question} sourceIds={response.sources || []} />
+          ) : null}
         </>
       )}
       {response.audio_en && (

src/components/App/UniverseQuestion/index.tsx

@@ -16,7 +16,7 @@ import { colors } from '~/utils/colors'
 export const UniverseQuestion = () => {
   const [question, setQuestion] = useState('')
   const textAreaRef = useRef<HTMLTextAreaElement>(null)
-  const { fetchData, setAbortRequests, seedQuestions } = useDataStore((s) => s)
+  const { setAbortRequests, seedQuestions, resetData } = useDataStore((s) => s)
   const [setBudget] = useUserStore((s) => [s.setBudget])
 
   const { setUniverseQuestionIsOpen, setSidebarOpen, setShowCollapseButton } = useAppStore((s) => ({
@@ -26,6 +26,7 @@ export const UniverseQuestion = () => {
   }))
 
   const resetAiSummaryAnswer = useAiSummaryStore((s) => s.resetAiSummaryAnswer)
+  const fetchAIData = useAiSummaryStore((s) => s.fetchAIData)
   const [displayedSeedQuestions, setDisplayedSeedQuestions] = useState<string[]>([])
 
   useEffect(() => {
@@ -44,11 +45,11 @@ export const UniverseQuestion = () => {
     if (questionToSubmit) {
       resetAiSummaryAnswer()
       setUniverseQuestionIsOpen()
+      resetData()
       setSidebarOpen(true)
       setShowCollapseButton(true)
+      await fetchAIData(setBudget, setAbortRequests, questionToSubmit)
     }
-
-    await fetchData(setBudget, setAbortRequests, questionToSubmit)
   }
 
   const canSubmit = !!question.trim().length

src/components/Universe/Graph/Connections/index.tsx

@@ -15,44 +15,53 @@ type Props = {
 export const Connections = memo(({ linksPosition }: Props) => {
   const [dataInitial, nodesNormalized] = useDataStore((s) => [s.dataInitial, s.nodesNormalized])
 
-  const [showSelectionGraph, searchQuery, selectedNodeTypes] = useGraphStore((s) => [
+  const [showSelectionGraph, searchQuery, selectedNodeTypes, highlightNodes, hoveredNode] = useGraphStore((s) => [
     s.showSelectionGraph,
     s.searchQuery,
     s.selectedNodeTypes,
+    s.highlightNodes,
+    s.hoveredNode,
   ])
 
   const normalizedSchemasByType = useSchemaStore((s) => s.normalizedSchemasByType)
 
   return (
     <>
       <group name="simulation-3d-group__connections" visible={!showSelectionGraph}>
-        {dataInitial?.links.map((l: Link) => {
-          const position = linksPosition.get(l.ref_id) || {
-            sx: 0,
-            sy: 0,
-            sz: 0,
-            tx: 0,
-            ty: 0,
-            tz: 0,
-          }
+        {dataInitial?.links?.length ? (
+          <>
+            {dataInitial?.links.map((l: Link) => {
+              const position = linksPosition.get(l.ref_id) || {
+                sx: 0,
+                sy: 0,
+                sz: 0,
+                tx: 0,
+                ty: 0,
+                tz: 0,
+              }
 
-          return (
-            <LineComponent
-              key={l.ref_id}
-              label={l.edge_type}
-              source={l.source}
-              sourceX={position.sx}
-              sourceY={position.sy}
-              sourceZ={position.sz}
-              target={l.target}
-              targetX={position.tx}
-              targetY={position.ty}
-              targetZ={position.tz}
-            />
-          )
-        })}
+              return (
+                <LineComponent
+                  key={l.ref_id}
+                  label={l.edge_type}
+                  source={l.source}
+                  sourceX={position.sx}
+                  sourceY={position.sy}
+                  sourceZ={position.sz}
+                  target={l.target}
+                  targetX={position.tx}
+                  targetY={position.ty}
+                  targetZ={position.tz}
+                />
+              )
+            })}
+          </>
+        ) : null}
       </group>
-      <group visible={!searchQuery && !selectedNodeTypes.length}>
+      <group
+        key={dataInitial?.links.length}
+        visible={!searchQuery && !selectedNodeTypes.length && !highlightNodes.length && !hoveredNode}
+      >
         <Segments limit={1000} lineWidth={0.05}>
           {dataInitial?.links.map((l: Link) => {
             const position = linksPosition.get(l.ref_id) || {

src/components/Universe/Graph/Cubes/SelectionDataNodes/Node/index.tsx

@@ -53,7 +53,7 @@ export const Node = ({ onClick, node, selected, rounded = true, x, y, z, id }: P
   // const iconName = Icon ? primaryIcon : 'NodesIcon'
   const keyProperty = getNodeKeysByType(node.node_type) || ''
 
-  const title = node?.properties ? node?.properties[keyProperty] : ''
+  const title = node?.properties && keyProperty ? node?.properties[keyProperty] || '' : node.name || ''
   const titleShortened = title ? truncateText(title, 30) : ''
   const description = keyProperty !== 'description' ? node.properties?.description : ''
   const descriptionShortened = description ? truncateText(description, 60) : ''

src/components/Universe/Graph/Cubes/SelectionDataNodes/index.tsx

@@ -26,6 +26,7 @@ export const SelectionDataNodes = memo(() => {
   const [selectionData, setSelectionData] = useState<GraphData | null>()
 
   const { addNewNode, nodesNormalized } = useDataStore((s) => s)
+  const dataInitial = useDataStore((s) => s.dataInitial)
   const selectedNode = useSelectedNode()
   const groupRef = useRef<Group>(null)
 
@@ -128,9 +129,41 @@ export const SelectionDataNodes = memo(() => {
     }
 
     if (selectedNode) {
-      init()
+      if (selectedNode.node_type === 'Question') {
+        const normalizedNode = nodesNormalized.get(selectedNode.ref_id)
+
+        if (!normalizedNode) {
+          return
+        }
+
+        const sourceNodes = (normalizedNode.sources || [])
+          .map((i) => nodesNormalized.get(i))
+          .filter((i) => !!i) as NodeExtended[]
+
+        const targetNodes = (normalizedNode.targets || [])
+          .map((i) => nodesNormalized.get(i))
+          .filter((i) => !!i) as NodeExtended[]
+
+        const siblings: NodeExtended[] = [
+          ...sourceNodes,
+          ...targetNodes,
+          { ...selectedNode, x: 0, y: 0, z: 0, fx: 0, fy: 0, fz: 0 },
+        ]
+
+        const links = (dataInitial?.links || []).filter((l: Link) =>
+          siblings.some(
+            (i: NodeExtended) =>
+              (i.ref_id === l.source && l.target === selectedNode.ref_id) ||
+              (i.ref_id === l.target && l.source === selectedNode.ref_id),
+          ),
+        )
+
+        setSelectionData({ nodes: siblings, links: links as unknown as GraphData['links'] })
+      } else {
+        init()
+      }
     }
-  }, [addNewNode, prevSelectedNodeId, selectedNode, setSelectionData])
+  }, [addNewNode, prevSelectedNodeId, selectedNode, setSelectionData, nodesNormalized, dataInitial?.links])
 
   const handleSelect = useCallback(
     (id: string) => {

src/components/Universe/Graph/Cubes/Text/index.tsx

@@ -58,24 +58,29 @@ export const TextNode = memo(
     const keyProperty = getNodeKeysByType(node.node_type) || ''
 
     const sanitizedNodeName =
-      keyProperty && node?.properties ? removeEmojis(String(node?.properties[keyProperty] || '')) : ''
+      keyProperty && node?.properties ? removeEmojis(String(node?.properties[keyProperty] || '')) : node.name || ''
 
     useFrame(({ camera, clock }) => {
       if (!nodeRef.current || !eventHandlerRef.current) {
         return
       }
 
-      const { selectedNode, hoveredNode, activeEdge, searchQuery, selectedNodeTypes, selectedLinkTypes } =
-        useGraphStore.getState()
+      const {
+        selectedNode,
+        hoveredNode,
+        activeEdge,
+        searchQuery,
+        selectedNodeTypes,
+        selectedLinkTypes,
+        hoveredNodeSiblings,
+      } = useGraphStore.getState()
 
       const checkDistance = () => {
         const nodePosition = nodePositionRef.current.setFromMatrixPosition(nodeRef.current!.matrixWorld)
 
         if (nodeRef.current) {
           nodeRef.current.visible = ignoreDistance ? true : nodePosition.distanceTo(camera.position) < 1500
         }
-
-        // Set visibility based on distance
       }
 
       if (searchQuery.length < 3 && !selectedNodeTypes.length && !selectedLinkTypes.length) {
@@ -94,6 +99,7 @@ export const TextNode = memo(
         activeEdge?.source === node.ref_id ||
         (searchQuery && sanitizedNodeName.toLowerCase().includes(searchQuery.toLowerCase())) ||
         selectedNodeTypes.includes(node.node_type) ||
+        hoveredNodeSiblings.includes(node.ref_id) ||
         node.edgeTypes?.some((i) => selectedLinkTypes.includes(i))
 
       if (isActive) {