File size: 4,544 Bytes
1587277
91b56ef
6093500
 
1587277
6093500
 
 
a5320ee
6093500
 
 
 
 
 
91b56ef
6093500
 
 
 
 
 
 
 
1587277
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a5320ee
1587277
a5320ee
6093500
 
1587277
 
 
 
 
 
 
 
 
 
 
 
 
 
6093500
1587277
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a5320ee
 
6093500
1587277
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
from typing import List

import instructor
from graphviz import Digraph
from pydantic import BaseModel, Field

from groq import Groq
import os

# Initialize with API key
client = Groq(api_key=os.getenv("GROQ_API_KEY"))

# Enable instructor patches for Groq client
client = instructor.from_groq(client)
"""
from openai import OpenAI
client = instructor.from_openai(
    OpenAI(
        base_url="http://localhost:11434/v1",
        api_key="ollama",
    ),
    mode=instructor.Mode.JSON,
)
"""
llm = 'llama-3.1-8b-instant' if os.getenv("GROQ_API_KEY") else "llama3.2"

class Node(BaseModel, frozen=True):
  """
  Node representing concept in the subject domain
  """
  id: int= Field(...,
                  description="unique id of the concept in the subject domain, used for deduplication, design a scheme allows multiple concept")
  label: str = Field(..., description="description of the concept in the subject domain")
  color: str = "orange"


class Edge(BaseModel, frozen=True):
  """
  Edge representing relationship between concepts in the subject domain, source depends on target
  """
  source: int = Field(..., description="source representing concept in the subject domain")
  target: int = Field(..., description="target representing concept in the subject domain")
  label: str = Field(..., description="description representing relationship between concepts in the subject domain")
  color: str = "black"


from typing import Optional


class KnowledgeGraph(BaseModel):
  """
    KnowledgeGraph is graph representation of concepts in the subject domain
    """
  nodes: Optional[List[Node]] = Field(..., default_factory=list)
  edges: Optional[List[Edge]] = Field(..., default_factory=list)

  def update(self, other: "KnowledgeGraph") -> "KnowledgeGraph":
    """Updates the current graph with the other graph, deduplicating nodes and edges."""
    return KnowledgeGraph(
      nodes=list(set(self.nodes + other.nodes)),
      edges=list(set(self.edges + other.edges)),
    )

  def draw(self, prefix: str = "knowledge_graph"):
    dot = Digraph(comment="Knowledge Graph")

    for node in self.nodes:
      dot.node(str(node.id), node.label, color=node.color)

    for edge in self.edges:
      dot.edge(
        str(edge.source), str(edge.target), label=edge.label, color=edge.color
      )
    dot.render(prefix, format="png", view=True)


from typing import Iterable
from textwrap import dedent


def generate_graph(q, input) -> KnowledgeGraph:
  return client.chat.completions.create(
    model=llm,
    max_retries=5,
    messages=[
      {
        "role": "user",
        "content": dedent(f"""Help me understand the following by describing it as a detailed knowledge graph: 
            ### Question: {q} 
            ### Context: {input}
            ### Instruction:
            Generate at least 5 concepts
            Generate at least 3 relationship
            ### Output Format:
            Node with id, label for description of the concept
            Edge with source's id, target's id, label for description of the relationship between source concept and target concept
            """),

      }
    ],
    response_model=KnowledgeGraph)


class Subissue(BaseModel):
  subissue_title: str
  point: List[str] = Field(default_factory=list, description="Specific aspect or component of the subissue")


def expandIssue(input) -> Iterable[Subissue]:
  response = client.chat.completions.create(
    model=llm,
    max_retries=3,
    response_model=Iterable[Subissue],
    temperature=0.1,
    messages=[

      {
        "role": "user",
        "content": dedent(f"""
                As a McKinsey Consultant, perform MECE decomposition of the question.
                ### Requirements
                1. Return 3 subissues minimum
                2. Each sub-issue has 3 bullet points, which each new point beginning with a *
                3. Use EXACT format:

                - [Sub-issue 1.1 title]
                  * [point 1]
                  * [point 2]
                  * [point 3]
                - [Sub-issue 1.2 title]
                  * [point 1]
                  * [point 2]
                  * [point 3]
                - [Sub-issue 1.3 title]
                  * [point 1]
                  * [point 2]
                  * [point 3]

                4. return nothing else
                ### Question: {input}
                """),
      },
    ],
  )

  return response


def graph(query):
  queryx = expandIssue(query)

  graph = generate_graph(query, str(queryx))
  return graph.json()